From 536e003c49f2ff1031c0bbf4f3a5507bfcfcf84a Mon Sep 17 00:00:00 2001
From: Ryo Kitagawa <kitadrum50@gmail.com>
Date: Sun, 15 Dec 2024 17:26:49 +0900
Subject: [PATCH] feat: support table without backtick

---
 langserver/internal/source/file/analyze.go   |  1 +
 langserver/internal/source/file/catalog.go   | 37 +++++++++++++-------
 langserver/internal/source/file/file_test.go | 14 ++++++++
 3 files changed, 40 insertions(+), 12 deletions(-)

diff --git a/langserver/internal/source/file/analyze.go b/langserver/internal/source/file/analyze.go
index 223485b..0d1fa39 100644
--- a/langserver/internal/source/file/analyze.go
+++ b/langserver/internal/source/file/analyze.go
@@ -40,6 +40,7 @@ func (a *Analyzer) langOpt() (*zetasql.LanguageOptions, error) {
 	langOpt.EnableLanguageFeature(zetasql.FeatureV13Qualify)
 	langOpt.EnableLanguageFeature(zetasql.FeatureV13ScriptLabel)
 	langOpt.EnableLanguageFeature(zetasql.FeatureAnalyticFunctions)
+	langOpt.EnableLanguageFeature(zetasql.FeatureV13AllowDashesInTableName)
 	langOpt.SetSupportsAllStatementKinds()
 	langOpt.EnableAllReservableKeywords(true)
 	err := langOpt.EnableReservableKeyword("QUALIFY", true)
diff --git a/langserver/internal/source/file/catalog.go b/langserver/internal/source/file/catalog.go
index e872ddf..cd48501 100644
--- a/langserver/internal/source/file/catalog.go
+++ b/langserver/internal/source/file/catalog.go
@@ -55,7 +55,14 @@ func (c *Catalog) FindTable(path []string) (types.Table, error) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
-	table, err := c.catalog.FindTable(path)
+	projectID, datasetID, tableID, err := c.pathToProjectTable(path)
+	if err != nil {
+		return nil, err
+	}
+
+	formattedPath := []string{fmt.Sprintf("%s.%s.%s", projectID, datasetID, tableID)}
+
+	table, err := c.catalog.FindTable(formattedPath)
 	if err == nil {
 		return table, nil
 	}
@@ -66,25 +73,21 @@ func (c *Catalog) FindTable(path []string) (types.Table, error) {
 		errs = append(errs, fmt.Errorf("failed to add table: %w", err))
 		return nil, errors.Join(errs...)
 	}
-	return c.catalog.FindTable(path)
+	return c.catalog.FindTable(formattedPath)
 }
 
 func (c *Catalog) addTable(path []string) error {
-	tableSep := strings.Split(strings.Join(path, "."), ".")
-	var metadata *bq.TableMetadata
-	var err error
-	if len(tableSep) == 3 {
-		metadata, err = c.bqClient.GetTableMetadata(context.Background(), tableSep[0], tableSep[1], tableSep[2])
-	} else if len(tableSep) == 2 {
-		metadata, err = c.bqClient.GetTableMetadata(context.Background(), c.bqClient.GetDefaultProject(), tableSep[0], tableSep[1])
-	} else {
-		return fmt.Errorf("unknown table: %s", strings.Join(path, "."))
+	projectID, datasetID, tableID, err := c.pathToProjectTable(path)
+	if err != nil {
+		return err
 	}
+
+	metadata, err := c.bqClient.GetTableMetadata(context.Background(), projectID, datasetID, tableID)
 	if err != nil {
 		return fmt.Errorf("failed to get schema: %w", err)
 	}
 
-	tableName := strings.Join(path, ".")
+	tableName := fmt.Sprintf("%s.%s.%s", projectID, datasetID, tableID)
 
 	schema := metadata.Schema
 	columns := make([]types.Column, len(schema))
@@ -110,6 +113,16 @@ func (c *Catalog) addTable(path []string) error {
 	return nil
 }
 
+func (c *Catalog) pathToProjectTable(path []string) (projectID, datasetID, tableID string, err error) {
+	tableSep := strings.Split(strings.Join(path, "."), ".")
+	if len(tableSep) == 3 {
+		return tableSep[0], tableSep[1], tableSep[2], nil
+	} else if len(tableSep) == 2 {
+		return c.bqClient.GetDefaultProject(), tableSep[0], tableSep[1], nil
+	}
+	return "", "", "", fmt.Errorf(`unknown table "%s"`, strings.Join(tableSep, "."))
+}
+
 func bigqueryTypeToZetaSQLType(typ bq.FieldType, isRepeated bool, schema bq.Schema) (types.Type, error) {
 	result, err := literalBigqueryTypeToZetaSQLType(typ, schema)
 	if err != nil {
diff --git a/langserver/internal/source/file/file_test.go b/langserver/internal/source/file/file_test.go
index f7abb3d..242eba2 100644
--- a/langserver/internal/source/file/file_test.go
+++ b/langserver/internal/source/file/file_test.go
@@ -584,6 +584,20 @@ func TestProject_ParseFile(t *testing.T) {
 			},
 			expectedErrs: []file.Error{},
 		},
+		"parse without backtick": {
+			file: "SELECT * FROM project.dataset.table",
+			bqTableMetadataMap: map[string]*bq.TableMetadata{
+				"project.dataset.table": {
+					Schema: bq.Schema{
+						{
+							Name: "id",
+							Type: bq.IntegerFieldType,
+						},
+					},
+				},
+			},
+			expectedErrs: []file.Error{},
+		},
 	}
 
 	for n, tt := range tests {