Completion column

langserver/completion.go

@@ -0,0 +1,31 @@
+package langserver
+
+import (
+	"context"
+	"encoding/json"
+
+	"github.com/kitagry/bqls/langserver/internal/lsp"
+	"github.com/sourcegraph/jsonrpc2"
+)
+
+func (h *Handler) handleTextDocumentCompletion(ctx context.Context, conn *jsonrpc2.Conn, req *jsonrpc2.Request) (result interface{}, err error) {
+	if req.Params == nil {
+		return nil, &jsonrpc2.Error{Code: jsonrpc2.CodeInvalidParams}
+	}
+
+	var params lsp.TextDocumentPositionParams
+	if err := json.Unmarshal(*req.Params, &params); err != nil {
+		return nil, err
+	}
+
+	items, err := h.project.Complete(ctx, documentURIToURI(params.TextDocument.URI), params.Position, h.clientSupportSnippets())
+	if err != nil {
+		return nil, err
+	}
+
+	return items, nil
+}
+
+func (h *Handler) clientSupportSnippets() bool {
+	return h.initializeParams.Capabilities.TextDocument.Completion.CompletionItem.SnippetSupport
+}

langserver/initialize.go

@@ -35,6 +35,10 @@ func (h *Handler) handleInitialize(ctx context.Context, conn *jsonrpc2.Conn, req
 			},
 			DocumentFormattingProvider: true,
 			HoverProvider:              true,
+			CompletionProvider: &lsp.CompletionOptions{
+				ResolveProvider:   true,
+				TriggerCharacters: []string{"*", "."},
+			},
 		},
 	}, nil
 }

langserver/internal/source/completion.go

@@ -0,0 +1,92 @@
+package source
+
+import (
+	"context"
+	"fmt"
+	"strings"
+
+	"cloud.google.com/go/bigquery"
+	"github.com/goccy/go-zetasql/ast"
+	"github.com/kitagry/bqls/langserver/internal/lsp"
+)
+
+func (p *Project) Complete(ctx context.Context, uri string, position lsp.Position, supportSunippet bool) ([]lsp.CompletionItem, error) {
+	result := make([]lsp.CompletionItem, 0)
+	sql := p.cache.Get(uri)
+
+	fromNodes := listAstNode[*ast.FromClauseNode](sql.Node)
+	for _, fromNode := range fromNodes {
+		tableExpr := fromNode.TableExpression()
+		if tableExpr == nil {
+			continue
+		}
+
+		switch tableExpr := tableExpr.(type) {
+		case *ast.TablePathExpressionNode:
+			tableMeta, err := p.getTableMetadataFromTablePathExpressionNode(ctx, tableExpr)
+			if err != nil {
+				return nil, fmt.Errorf("failed to get table metadata: %w", err)
+			}
+
+			for _, schema := range tableMeta.Schema {
+				detail := string(schema.Type)
+				if schema.Description != "" {
+					detail += "\n" + schema.Description
+				}
+				if !supportSunippet {
+					result = append(result, lsp.CompletionItem{
+						InsertTextFormat: lsp.ITFPlainText,
+						Kind:             lsp.CIKField,
+						Label:            schema.Name,
+						Detail:           detail,
+					})
+					continue
+				}
+
+				result = append(result, lsp.CompletionItem{
+					InsertTextFormat: lsp.ITFSnippet,
+					Kind:             lsp.CIKField,
+					Label:            schema.Name,
+					Detail:           detail,
+					TextEdit: &lsp.TextEdit{
+						NewText: schema.Name,
+						Range: lsp.Range{
+							Start: position,
+							End:   position,
+						},
+					},
+				})
+			}
+		}
+	}
+
+	return result, nil
+}
+
+func (p *Project) getTableMetadataFromTablePathExpressionNode(ctx context.Context, tableNode *ast.TablePathExpressionNode) (*bigquery.TableMetadata, error) {
+	pathExpr := tableNode.PathExpr()
+	if pathExpr == nil {
+		return nil, fmt.Errorf("invalid table path expression")
+	}
+	pathNames := make([]string, len(pathExpr.Names()))
+	for i, n := range pathExpr.Names() {
+		pathNames[i] = n.Name()
+	}
+	targetTable, err := p.getTableMetadataFromPath(ctx, strings.Join(pathNames, "."))
+	if err != nil {
+		return nil, fmt.Errorf("failed to get table metadata: %w", err)
+	}
+	return targetTable, nil
+}
+
+func listAstNode[T ast.Node](node ast.Node) []T {
+	targetNodes := make([]T, 0)
+	ast.Walk(node, func(n ast.Node) error {
+		node, ok := n.(T)
+		if ok {
+			targetNodes = append(targetNodes, node)
+		}
+		return nil
+	})
+	return targetNodes
+}

langserver/internal/source/completion_test.go

@@ -0,0 +1,134 @@
+package source_test
+
+import (
+	"context"
+	"errors"
+	"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"
+	"github.com/kitagry/bqls/langserver/internal/source/helper"
+)
+
+func TestProject_Complete(t *testing.T) {
+	tests := map[string]struct {
+		files           map[string]string
+		supportSunippet bool
+		bqTableMetadata *bq.TableMetadata
+
+		expectCompletionItems []lsp.CompletionItem
+		expectErr             error
+	}{
+		"Select columns with supportSunippet is false": {
+			files: map[string]string{
+				"file1.sql": "SELECT id, | FROM `project.dataset.table`",
+			},
+			supportSunippet: false,
+			bqTableMetadata: &bq.TableMetadata{
+				Schema: bq.Schema{
+					{
+						Name:        "id",
+						Type:        bq.IntegerFieldType,
+						Description: "id description",
+					},
+					{
+						Name: "name",
+						Type: bq.StringFieldType,
+					},
+				},
+			},
+			expectCompletionItems: []lsp.CompletionItem{
+				{
+					InsertTextFormat: lsp.ITFPlainText,
+					Kind:             lsp.CIKField,
+					Label:            "id",
+					Detail:           "INTEGER\nid description",
+				},
+				{
+					InsertTextFormat: lsp.ITFPlainText,
+					Kind:             lsp.CIKField,
+					Label:            "name",
+					Detail:           "STRING",
+				},
+			},
+		},
+		"Select columns with supportSunippet is true": {
+			files: map[string]string{
+				"file1.sql": "SELECT id, | FROM `project.dataset.table`",
+			},
+			supportSunippet: true,
+			bqTableMetadata: &bq.TableMetadata{
+				Schema: bq.Schema{
+					{
+						Name:        "id",
+						Type:        bq.IntegerFieldType,
+						Description: "id description",
+					},
+					{
+						Name: "name",
+						Type: bq.StringFieldType,
+					},
+				},
+			},
+			expectCompletionItems: []lsp.CompletionItem{
+				{
+					InsertTextFormat: lsp.ITFSnippet,
+					Kind:             lsp.CIKField,
+					Label:            "id",
+					Detail:           "INTEGER\nid description",
+					TextEdit: &lsp.TextEdit{
+						NewText: "id",
+						Range: lsp.Range{
+							Start: lsp.Position{Line: 0, Character: 11},
+							End:   lsp.Position{Line: 0, Character: 11},
+						},
+					},
+				},
+				{
+					InsertTextFormat: lsp.ITFSnippet,
+					Kind:             lsp.CIKField,
+					Label:            "name",
+					Detail:           "STRING",
+					TextEdit: &lsp.TextEdit{
+						NewText: "id",
+						Range: lsp.Range{
+							Start: lsp.Position{Line: 0, Character: 11},
+							End:   lsp.Position{Line: 0, Character: 11},
+						},
+					},
+				},
+			},
+		},
+	}
+
+	for n, tt := range tests {
+		t.Run(n, func(t *testing.T) {
+			ctrl := gomock.NewController(t)
+			bqClient := mock_bigquery.NewMockClient(ctrl)
+			bqClient.EXPECT().GetTableMetadata(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(tt.bqTableMetadata, nil).MinTimes(0)
+			p := source.NewProjectWithBQClient("/", bqClient)
+
+			files, path, position, err := helper.GetLspPosition(tt.files)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			for uri, content := range files {
+				p.UpdateFile(uri, content, 1)
+			}
+
+			got, err := p.Complete(context.Background(), path, position, tt.supportSunippet)
+			if !errors.Is(err, tt.expectErr) {
+				t.Fatalf("got error %v, but want %v", err, tt.expectErr)
+			}
+
+			if diff := cmp.Diff(got, tt.expectCompletionItems); diff != "" {
+				t.Errorf("(-got, +want)\n%s", diff)
+			}
+		})
+	}
+}

langserver/langserver.go

@@ -63,6 +63,8 @@ func (h *Handler) handle(ctx context.Context, conn *jsonrpc2.Conn, req *jsonrpc2
 		return h.handleTextDocumentFormatting(ctx, conn, req)
 	case "textDocument/hover":
 		return h.handleTextDocumentHover(ctx, conn, req)
+	case "textDocument/completion":
+		return h.handleTextDocumentCompletion(ctx, conn, req)
 	}
 	return nil, &jsonrpc2.Error{Code: jsonrpc2.CodeMethodNotFound, Message: fmt.Sprintf("method not supported: %s", req.Method)}
 }