Merge pull request #124 from kitagry/improve-list-table-performance

improve list table performance by removing onlyLatestSuffix
kitagry · Mar 2, 2024 · 0196475 · 0196475
2 parents aba5c5c + 7ffb775
commit 0196475
Show file tree

Hide file tree

Showing 9 changed files with 48 additions and 27 deletions.
diff --git a/langserver/internal/bigquery/bigquery.go b/langserver/internal/bigquery/bigquery.go
@@ -3,6 +3,7 @@ package bigquery
 import (
 	"context"
 	"fmt"
+	"strings"
 
 	"cloud.google.com/go/bigquery"
 	"google.golang.org/api/cloudresourcemanager/v1"
@@ -22,7 +23,7 @@ type Client interface {
 	ListDatasets(ctx context.Context, projectID string) ([]*bigquery.Dataset, error)
 
 	// ListTables lists all tables in the specified dataset.
-	ListTables(ctx context.Context, projectID, datasetID string, onlyLatestSuffix bool) ([]*bigquery.Table, error)
+	ListTables(ctx context.Context, projectID, datasetID string) ([]*bigquery.Table, error)
 
 	// GetTableMetadata returns the metadata of the specified table.
 	GetTableMetadata(ctx context.Context, projectID, datasetID, tableID string) (*bigquery.TableMetadata, error)
@@ -117,10 +118,11 @@ func (c *client) ListDatasets(ctx context.Context, projectID string) ([]*bigquer
 	return datasets, nil
 }
 
-func (c *client) ListTables(ctx context.Context, projectID, datasetID string, onlyLatestSuffix bool) ([]*bigquery.Table, error) {
+func (c *client) ListTables(ctx context.Context, projectID, datasetID string) ([]*bigquery.Table, error) {
 	dataset := c.bqClient.DatasetInProject(projectID, datasetID)
 
 	it := dataset.Tables(ctx)
+	it.PageInfo().MaxSize = 1000
 
 	tables := make([]*bigquery.Table, 0)
 	for {
@@ -132,12 +134,15 @@ func (c *client) ListTables(ctx context.Context, projectID, datasetID string, on
 			return nil, fmt.Errorf("fail to scan DatasetInProject: %w", err)
 		}
 
+		if strings.HasPrefix(table.TableID, "LOAD_TEMP_") {
+			continue
+		}
+
 		tables = append(tables, table)
 	}
 
-	if onlyLatestSuffix {
-		tables = extractLatestSuffixTables(tables)
-	}
+	tables = extractLatestSuffixTables(tables)
+
 	return tables, nil
 }
 

diff --git a/langserver/internal/bigquery/cache.go b/langserver/internal/bigquery/cache.go
@@ -129,7 +129,7 @@ func (c *cache) callListDatasets(ctx context.Context, projectID string) ([]*bigq
 	return result, nil
 }
 
-func (c *cache) ListTables(ctx context.Context, projectID, datasetID string, onlyLatestSuffix bool) ([]*bigquery.Table, error) {
+func (c *cache) ListTables(ctx context.Context, projectID, datasetID string) ([]*bigquery.Table, error) {
 	results, err := c.db.SelectTables(ctx, projectID, datasetID)
 	if err == nil && len(results) > 0 {
 		key := fmt.Sprintf("%s.%s", projectID, datasetID)
@@ -139,28 +139,25 @@ func (c *cache) ListTables(ctx context.Context, projectID, datasetID string, onl
 
 		go c.onceListTables[key].Do(func() {
 			ctx := context.WithoutCancel(ctx)
-			_, err := c.callListTables(ctx, projectID, datasetID, false)
+			_, err := c.callListTables(ctx, projectID, datasetID)
 			if err != nil {
 				fmt.Fprintf(os.Stderr, "failed to recache tables: %v\n", err)
 			}
 		})
 
-		if onlyLatestSuffix {
-			return extractLatestSuffixTables(results), nil
-		}
 		return results, nil
 	}
 	if err != nil {
 		// TODO
 		fmt.Fprintf(os.Stderr, "failed to select tables: %v\n", err)
 	}
 
-	return c.callListTables(ctx, projectID, datasetID, onlyLatestSuffix)
+	return c.callListTables(ctx, projectID, datasetID)
 }
 
-func (c *cache) callListTables(ctx context.Context, projectID, datasetID string, onlyLatestSuffix bool) ([]*bigquery.Table, error) {
+func (c *cache) callListTables(ctx context.Context, projectID, datasetID string) ([]*bigquery.Table, error) {
 	// onlyLatestSuffix is ignored for cache.
-	result, err := c.bqClient.ListTables(ctx, projectID, datasetID, false)
+	result, err := c.bqClient.ListTables(ctx, projectID, datasetID)
 	if err != nil {
 		return nil, err
 	}
@@ -173,9 +170,6 @@ func (c *cache) callListTables(ctx context.Context, projectID, datasetID string,
 		}
 	}
 
-	if onlyLatestSuffix {
-		result = extractLatestSuffixTables(result)
-	}
 	return result, nil
 }
 

diff --git a/langserver/internal/bigquery/mock_bigquery/mock_bigquery.go b/langserver/internal/bigquery/mock_bigquery/mock_bigquery.go
diff --git a/langserver/internal/function/builtin_function.go b/langserver/internal/function/builtin_function.go
diff --git a/langserver/internal/source/completion/column_test.go b/langserver/internal/source/completion/column_test.go
@@ -822,7 +822,7 @@ func TestProject_CompleteColumns(t *testing.T) {
 				}
 				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)
+			bqClient.EXPECT().ListTables(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).MinTimes(0)
 			logger := logrus.New()
 			logger.SetLevel(logrus.DebugLevel)
 

diff --git a/langserver/internal/source/completion/declaration_test.go b/langserver/internal/source/completion/declaration_test.go
@@ -65,7 +65,7 @@ func TestProject_CompleteDeclaration(t *testing.T) {
 				}
 				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)
+			bqClient.EXPECT().ListTables(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).MinTimes(0)
 			logger := logrus.New()
 			logger.SetLevel(logrus.DebugLevel)
 

diff --git a/langserver/internal/source/completion/table_path.go b/langserver/internal/source/completion/table_path.go
@@ -107,7 +107,7 @@ func (c *completor) completeDatasetForTablePath(ctx context.Context, param table
 }
 
 func (c *completor) completeTableForTablePath(ctx context.Context, param tablePathParams) ([]CompletionItem, error) {
-	tables, err := c.bqClient.ListTables(ctx, param.ProjectID, param.DatasetID, true)
+	tables, err := c.bqClient.ListTables(ctx, param.ProjectID, param.DatasetID)
 	if err != nil {
 		return nil, fmt.Errorf("failed to bqClient.ListTables: %w", err)
 	}

diff --git a/langserver/internal/source/completion/table_path_test.go b/langserver/internal/source/completion/table_path_test.go
@@ -34,7 +34,7 @@ func TestProject_CompleteTablePath(t *testing.T) {
 				ctrl := gomock.NewController(t)
 				bqClient := mock_bigquery.NewMockClient(ctrl)
 
-				bqClient.EXPECT().ListTables(gomock.Any(), "project", "dataset", gomock.Any()).Return([]*bq.Table{
+				bqClient.EXPECT().ListTables(gomock.Any(), "project", "dataset").Return([]*bq.Table{
 					{
 						ProjectID: "project",
 						DatasetID: "dataset",
@@ -76,7 +76,7 @@ func TestProject_CompleteTablePath(t *testing.T) {
 				ctrl := gomock.NewController(t)
 				bqClient := mock_bigquery.NewMockClient(ctrl)
 
-				bqClient.EXPECT().ListTables(gomock.Any(), "project", "dataset", true).Return([]*bq.Table{
+				bqClient.EXPECT().ListTables(gomock.Any(), "project", "dataset").Return([]*bq.Table{
 					{
 						ProjectID: "project",
 						DatasetID: "dataset",

diff --git a/langserver/internal/source/project.go b/langserver/internal/source/project.go
@@ -146,7 +146,7 @@ func (p *Project) ListDatasets(ctx context.Context, projectID string) ([]*bq.Dat
 }
 
 func (p *Project) ListTables(ctx context.Context, projectID, datasetID string) ([]*bq.Table, error) {
-	return p.bqClient.ListTables(ctx, projectID, datasetID, true)
+	return p.bqClient.ListTables(ctx, projectID, datasetID)
 }
 
 func (p *Project) ListJobs(ctx context.Context, projectID string, allUsers bool) ([]lsp.JobHistory, error) {