Merge pull request #78 from kitagry/support-time-partitioning

Support _PARTITIONTIME column
kitagry · Oct 26, 2023 · 9ff32f0 · 9ff32f0
2 parents a453475 + c5ea0dd
commit 9ff32f0
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/langserver/internal/source/file/catalog.go b/langserver/internal/source/file/catalog.go
@@ -83,6 +83,10 @@ func (c *Catalog) addTable(path []string) error {
 		columns[i] = types.NewSimpleColumn(tableName, field.Name, typ)
 	}
 
+	if metadata.TimePartitioning != nil {
+		columns = append(columns, types.NewSimpleColumn(tableName, "_PARTITIONTIME", types.TimestampType()))
+	}
+
 	if isWildCardTable(path) {
 		columns = append(columns, types.NewSimpleColumn(tableName, "_TABLE_SUFFIX", types.StringType()))
 	}

diff --git a/langserver/internal/source/file/catalog_test.go b/langserver/internal/source/file/catalog_test.go
@@ -67,6 +67,26 @@ func TestCatalog_AddTable(t *testing.T) {
 			expectTableName:   "project.dataset.table*",
 			expectColumnNames: []string{"name", "_TABLE_SUFFIX"},
 		},
+		"partitiontime table": {
+			path: []string{"project.dataset.table"},
+			createMockBigQuery: func(ctrl *gomock.Controller) bigquery.Client {
+				bqClient := mock_bigquery.NewMockClient(ctrl)
+				bqClient.EXPECT().GetTableMetadata(gomock.Any(), "project", "dataset", "table").Return(&bq.TableMetadata{
+					Schema: bq.Schema{
+						{
+							Name: "name",
+							Type: bq.StringFieldType,
+						},
+					},
+					TimePartitioning: &bq.TimePartitioning{
+						Type: bq.DayPartitioningType,
+					},
+				}, nil)
+				return bqClient
+			},
+			expectTableName:   "project.dataset.table",
+			expectColumnNames: []string{"name", "_PARTITIONTIME"},
+		},
 	}
 
 	for n, tt := range tests {