Merge pull request #59 from kitagry/complete-declaration

Complete declaration

Author: kitagry

langserver/internal/source/completion/completion.go

@@ -35,6 +35,7 @@ func (c *completor) Complete(ctx context.Context, parsedFile file.ParsedFile, po
 
 	result = append(result, c.completeColumns(ctx, parsedFile, position)...)
 	result = append(result, c.completeBuiltinFunction(ctx, parsedFile, position)...)
+	result = append(result, c.completeDeclaration(ctx, parsedFile, position)...)
 	return result, nil
 }
 

langserver/internal/source/completion/declaration.go

@@ -0,0 +1,45 @@
+package completion
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"github.com/goccy/go-zetasql"
+	"github.com/goccy/go-zetasql/ast"
+	"github.com/kitagry/bqls/langserver/internal/lsp"
+	"github.com/kitagry/bqls/langserver/internal/source/file"
+)
+
+func (c *completor) completeDeclaration(ctx context.Context, parsedFile file.ParsedFile, position lsp.Position) []CompletionItem {
+	termOffset := parsedFile.TermOffset(position)
+
+	// When the cursor is in the middle of the table path, do not suggest the built-in functions.
+	tablePathNode, ok := file.SearchAstNode[*ast.TablePathExpressionNode](parsedFile.Node, parsedFile.TermOffset(position))
+	if ok && tablePathNode.ParseLocationRange().End().ByteOffset() != termOffset {
+		return []CompletionItem{}
+	}
+
+	incompleteColumnName := parsedFile.FindIncompleteColumnName(position)
+
+	declarations := file.ListAstNode[*ast.VariableDeclarationNode](parsedFile.Node)
+
+	result := make([]CompletionItem, 0, len(declarations))
+	for _, d := range declarations {
+		iList := d.VariableList().IdentifierList()
+		for _, l := range iList {
+			if !strings.HasPrefix(l.Name(), incompleteColumnName) {
+				continue
+			}
+			result = append(result, CompletionItem{
+				Kind:    lsp.CIKVariable,
+				NewText: l.Name(),
+				Documentation: lsp.MarkupContent{
+					Kind:  lsp.MKMarkdown,
+					Value: fmt.Sprintf("```sql\n%s```", zetasql.Unparse(d)),
+				},
+			})
+		}
+	}
+	return result
+}

langserver/internal/source/completion/declaration_test.go

@@ -0,0 +1,88 @@
+package completion
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	bq "cloud.google.com/go/bigquery"
+	"github.com/golang/mock/gomock"
+	"github.com/google/go-cmp/cmp"
+	"github.com/kitagry/bqls/langserver/internal/bigquery/mock_bigquery"
+	"github.com/kitagry/bqls/langserver/internal/lsp"
+	"github.com/kitagry/bqls/langserver/internal/source/file"
+	"github.com/kitagry/bqls/langserver/internal/source/helper"
+	"github.com/sirupsen/logrus"
+)
+
+func TestProject_CompleteDeclaration(t *testing.T) {
+	tests := map[string]struct {
+		files              map[string]string
+		bqTableMetadataMap map[string]*bq.TableMetadata
+
+		expectCompletionItems []CompletionItem
+	}{
+		"Complete variable declaration": {
+			files: map[string]string{
+				"file1.sql": "DECLARE a INT64;\n" + "SELECT | FROM `project.dataset.table`",
+			},
+			bqTableMetadataMap: map[string]*bq.TableMetadata{
+				"project.dataset.table": {
+					Schema: bq.Schema{
+						{
+							Name:        "id",
+							Type:        bq.IntegerFieldType,
+							Description: "id description",
+						},
+						{
+							Name: "name",
+							Type: bq.StringFieldType,
+						},
+					},
+				},
+			},
+			expectCompletionItems: []CompletionItem{
+				{
+					Kind:    lsp.CIKVariable,
+					NewText: "a",
+					Documentation: lsp.MarkupContent{
+						Kind:  lsp.MKMarkdown,
+						Value: "```sql\nDECLARE a INT64\n```",
+					},
+				},
+			},
+		},
+	}
+
+	for n, tt := range tests {
+		t.Run(n, func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			bqClient := mock_bigquery.NewMockClient(ctrl)
+			for tablePath, schema := range tt.bqTableMetadataMap {
+				tablePathSplitted := strings.Split(tablePath, ".")
+				if len(tablePathSplitted) != 3 {
+					t.Fatalf("table path length should be 3, got %s", tablePath)
+				}
+				bqClient.EXPECT().GetTableMetadata(gomock.Any(), tablePathSplitted[0], tablePathSplitted[1], tablePathSplitted[2]).Return(schema, nil).MinTimes(0)
+			}
+			bqClient.EXPECT().ListTables(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).MinTimes(0)
+			logger := logrus.New()
+			logger.SetLevel(logrus.DebugLevel)
+
+			analyzer := file.NewAnalyzer(logger, bqClient)
+			completor := New(logger, analyzer, bqClient)
+
+			files, path, position, err := helper.GetLspPosition(tt.files)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			parsedFile := analyzer.ParseFile(path, files[path])
+
+			got := completor.completeDeclaration(context.Background(), parsedFile, position)
+			if diff := cmp.Diff(got, tt.expectCompletionItems); diff != "" {
+				t.Errorf("(-got, +want)\n%s", diff)
+			}
+		})
+	}
+}

langserver/internal/source/file/analyze.go

@@ -246,14 +246,47 @@ func getDummyValueForDeclarationNode(node *ast.VariableDeclarationNode) (string,
 	case *ast.ArrayTypeNode:
 		return "[]", nil
 	case *ast.SimpleTypeNode:
-		return getDummyValueForDefaultValueNode(node.DefaultValue())
+		if node.DefaultValue() != nil {
+			return getDummyValueForDefaultValueNode(node.DefaultValue())
+		}
+		if pen, ok := n.Child(0).(*ast.PathExpressionNode); ok {
+			if in, ok := pen.Child(0).(*ast.IdentifierNode); ok {
+				return getDummyValueForDeclarationIdentifierName(in.Name())
+			}
+		}
+		return "", fmt.Errorf("failed to load default value")
 	default:
 		return "", fmt.Errorf("not implemented: %s", n.Kind())
 	}
 }
 
+func getDummyValueForDeclarationIdentifierName(name string) (string, error) {
+	switch name {
+	case "BOOL":
+		return "TRUE", nil
+	case "INT64":
+		return "1", nil
+	case "FLOAT64":
+		return "1.0", nil
+	case "STRING":
+		return "''", nil
+	case "BYTES":
+		return "b''", nil
+	case "DATE":
+		return "DATE('1970-01-01')", nil
+	case "DATETIME":
+		return "DATETIME('1970-01-01 00:00:00')", nil
+	case "TIME":
+		return "TIME('00:00:00')", nil
+	case "TIMESTAMP":
+		return "TIMESTAMP('1970-01-01 00:00:00')", nil
+	default:
+		return "", fmt.Errorf("not implemented: %s", name)
+	}
+}
+
 func getDummyValueForDefaultValueNode(node ast.ExpressionNode) (string, error) {
-	switch n := node.(type) {
+	switch node.(type) {
 	case *ast.NullLiteralNode:
 		return "NULL", nil
 	case *ast.BooleanLiteralNode:
@@ -267,6 +300,6 @@ func getDummyValueForDefaultValueNode(node ast.ExpressionNode) (string, error) {
 	case *ast.DateOrTimeLiteralNode:
 		return "DATE('1970-01-01')", nil
 	default:
-		return "", fmt.Errorf("not implemented: %s", n.Kind())
+		return "", fmt.Errorf("not implemented: %T", node)
 	}
 }

langserver/internal/source/file/file_test.go

@@ -646,7 +646,7 @@ func TestAnalyzer_ParseFileWithDeclareStatement(t *testing.T) {
 
 		expectedErrs []file.Error
 	}{
-		"Parse with incomplete table name": {
+		"Parse with default value": {
 			file: "DECLARE target_id INT64 DEFAULT 10;\n" +
 				"SELECT * FROM `project.dataset.table` WHERE id = target_id",
 			bqTableMetadataMap: map[string]*bq.TableMetadata{
@@ -661,6 +661,21 @@ func TestAnalyzer_ParseFileWithDeclareStatement(t *testing.T) {
 			},
 			expectedErrs: []file.Error{},
 		},
+		"Parse without default value": {
+			file: "DECLARE target_id INT64;\n" +
+				"SELECT * FROM `project.dataset.table` WHERE id = target_id",
+			bqTableMetadataMap: map[string]*bq.TableMetadata{
+				"project.dataset.table": {
+					Schema: bq.Schema{
+						{
+							Name: "id",
+							Type: bq.IntegerFieldType,
+						},
+					},
+				},
+			},
+			expectedErrs: []file.Error{},
+		},
 	}
 
 	for n, tt := range tests {

langserver/internal/source/file/node.go

@@ -91,3 +91,17 @@ func LookupNode[T astNode](n ast.Node) (T, bool) {
 
 	return LookupNode[T](n.Parent())
 }
+
+func ListAstNode[T locationRangeNode](n ast.ScriptNode) []T {
+	result := make([]T, 0)
+	ast.Walk(n, func(n ast.Node) error {
+		node, ok := n.(T)
+		if !ok {
+			return nil
+		}
+		result = append(result, node)
+		return nil
+	})
+
+	return result
+}