Merge pull request #199 from kitagry/update-no-project

feat: add dataset and table completion for default projectID
kitagry · Dec 26, 2024 · 5957efc · 5957efc
2 parents de44475 + 6a05f61
commit 5957efc
Show file tree

Hide file tree

Showing 4 changed files with 301 additions and 150 deletions.
diff --git a/langserver/internal/source/completion/table_path.go b/langserver/internal/source/completion/table_path.go
@@ -2,6 +2,7 @@ package completion
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"strings"
 
@@ -31,6 +32,8 @@ func (c *completor) completeTablePath(ctx context.Context, parsedFile file.Parse
 	}
 
 	splittedTablePath := strings.Split(tablePath, ".")
+
+	// for `project.dataset.table` completion
 	params := tablePathParams{}
 	if len(splittedTablePath) >= 1 {
 		params.ProjectID = splittedTablePath[0]
@@ -42,11 +45,32 @@ func (c *completor) completeTablePath(ctx context.Context, parsedFile file.Parse
 		params.TableID = splittedTablePath[2]
 	}
 
+	// for `dataset.table` completion
+	defaultParams := tablePathParams{
+		ProjectID: c.bqClient.GetDefaultProject(),
+	}
+	if len(splittedTablePath) >= 1 {
+		defaultParams.DatasetID = splittedTablePath[0]
+	}
+	if len(splittedTablePath) >= 2 {
+		defaultParams.TableID = splittedTablePath[1]
+	}
+
 	switch len(splittedTablePath) {
 	case 0, 1:
-		return c.completeProjectForTablePath(ctx, params)
+		result1, err1 := c.completeProjectForTablePath(ctx, params)
+		result2, err2 := c.completeDatasetForTablePath(ctx, defaultParams)
+		if len(result1) == 0 && len(result2) == 0 {
+			return nil, errors.Join(err1, err2)
+		}
+		return append(result1, result2...), nil
 	case 2:
-		return c.completeDatasetForTablePath(ctx, params)
+		result1, err1 := c.completeDatasetForTablePath(ctx, params)
+		result2, err2 := c.completeTableForTablePath(ctx, defaultParams)
+		if len(result1) == 0 && len(result2) == 0 {
+			return nil, errors.Join(err1, err2)
+		}
+		return append(result1, result2...), nil
 	case 3:
 		return c.completeTableForTablePath(ctx, params)
 	}

diff --git a/langserver/internal/source/completion/table_path_test.go b/langserver/internal/source/completion/table_path_test.go
@@ -34,6 +34,7 @@ func TestProject_CompleteTablePath(t *testing.T) {
 				ctrl := gomock.NewController(t)
 				bqClient := mock_bigquery.NewMockClient(ctrl)
 
+				bqClient.EXPECT().GetDefaultProject().Return("project")
 				bqClient.EXPECT().ListTables(gomock.Any(), "project", "dataset").Return([]*bq.Table{
 					{
 						ProjectID: "project",
@@ -76,6 +77,7 @@ func TestProject_CompleteTablePath(t *testing.T) {
 				ctrl := gomock.NewController(t)
 				bqClient := mock_bigquery.NewMockClient(ctrl)
 
+				bqClient.EXPECT().GetDefaultProject().Return("project")
 				bqClient.EXPECT().ListTables(gomock.Any(), "project", "dataset").Return([]*bq.Table{
 					{
 						ProjectID: "project",
@@ -116,7 +118,8 @@ func TestProject_CompleteTablePath(t *testing.T) {
 						DatasetID: "dataset2",
 					},
 				}, nil)
-				bqClient.EXPECT().GetDefaultProject().Return("").MinTimes(0)
+				bqClient.EXPECT().ListTables(gomock.Any(), "default_project", "project").Return([]*bq.Table{}, nil)
+				bqClient.EXPECT().GetDefaultProject().Return("default_project").MinTimes(0)
 				bqClient.EXPECT().GetTableMetadata(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, fmt.Errorf("not found")).MinTimes(0)
 				return bqClient
 			},
@@ -157,6 +160,8 @@ func TestProject_CompleteTablePath(t *testing.T) {
 						Name:      "project name",
 					},
 				}, nil)
+				bqClient.EXPECT().GetDefaultProject().Return("default_project")
+				bqClient.EXPECT().ListDatasets(gomock.Any(), "default_project").Return([]*bq.Dataset{}, nil)
 				bqClient.EXPECT().GetTableMetadata(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, fmt.Errorf("not found")).MinTimes(0)
 				return bqClient
 			},

diff --git a/langserver/internal/source/document.go b/langserver/internal/source/document.go
@@ -333,6 +333,9 @@ func (p *Project) findInputScan(name string, scanNode rast.ScanNode) (rast.ScanN
 			if n.Table().Name() == name {
 				return n, true
 			}
+			if n.Table().Name() == strings.Join([]string{p.bqClient.GetDefaultProject(), name}, ".") {
+				return n, true
+			}
 			return nil, false
 		case *rast.WithRefScanNode:
 			if n.WithQueryName() == name {