[zetasql-toolkit] fix: ensure BigQueryCatalog.addAllTablesUsedInQuery…

… ignores unqualified tables (GoogleCloudPlatform#1016)
sunnyly2016 · Apr 21, 2023 · c950001 · c950001
1 parent 3f33a54
commit c950001
Show file tree

Hide file tree

Showing 5 changed files with 43 additions and 24 deletions.
diff --git a/tools/zetasql-helper/zetasql-toolkit-core/pom.xml b/tools/zetasql-helper/zetasql-toolkit-core/pom.xml
@@ -6,7 +6,7 @@
 
     <groupId>com.google.zetasql.toolkit</groupId>
     <artifactId>zetasql-toolkit-core</artifactId>
-    <version>0.2.0</version>
+    <version>0.2.1</version>
 
     <properties>
         <zetasql.toolkit.version>${project.version}</zetasql.toolkit.version>

diff --git a/...zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/CatalogWrapper.java b/...zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/CatalogWrapper.java
@@ -16,8 +16,6 @@
 
 package com.google.zetasql.toolkit.catalog;
 
-import com.google.zetasql.Analyzer;
-import com.google.zetasql.AnalyzerOptions;
 import com.google.zetasql.SimpleCatalog;
 import com.google.zetasql.SimpleTable;
 import com.google.zetasql.resolvedast.ResolvedCreateStatementEnums.CreateMode;
@@ -26,8 +24,6 @@
 import com.google.zetasql.toolkit.catalog.bigquery.ProcedureInfo;
 import com.google.zetasql.toolkit.catalog.bigquery.TVFInfo;
 import java.util.List;
-import java.util.Set;
-import java.util.stream.Collectors;
 
 /**
  * Interface for an object that wraps a ZetaSQL SimpleCatalog and allows adding resources to it by
@@ -209,24 +205,6 @@ default void addProcedure(String procedure) {
     this.addProcedures(List.of(procedure));
   }
 
-  /**
-   * Adds all the tables used in the provided query to this catalog.
-   *
-   * <p>Uses Analyzer.extractTableNamesFromScript to extract the table names and later uses
-   * this.addTables to add them.
-   *
-   * @param query The SQL query from which to get the tables that should be added to the catalog
-   * @param options The ZetaSQL AnalyzerOptions to use when extracting the table names from the
-   *     query
-   */
-  default void addAllTablesUsedInQuery(String query, AnalyzerOptions options) {
-    Set<String> tables =
-        Analyzer.extractTableNamesFromScript(query, options).stream()
-            .map(tablePath -> String.join(".", tablePath))
-            .collect(Collectors.toSet());
-    this.addTables(List.copyOf(tables));
-  }
-
   /**
    * Creates a copy of this CatalogWrapper.
    *

diff --git a/...olkit-core/src/main/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryCatalog.java b/...olkit-core/src/main/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryCatalog.java
@@ -460,6 +460,25 @@ public void addAllTablesInProject(String projectId) {
                     table, CreateMode.CREATE_OR_REPLACE, CreateScope.CREATE_DEFAULT_SCOPE));
   }
 
+  /**
+   * Adds all the tables used in the provided query to this catalog.
+   *
+   * <p>Uses Analyzer.extractTableNamesFromScript to extract the table names and later uses
+   * this.addTables to add them.
+   *
+   * @param query The SQL query from which to get the tables that should be added to the catalog
+   * @param options The ZetaSQL AnalyzerOptions to use when extracting the table names from the
+   *     query
+   */
+  public void addAllTablesUsedInQuery(String query, AnalyzerOptions options) {
+    Set<String> tables =
+        Analyzer.extractTableNamesFromScript(query, options).stream()
+            .map(tablePath -> String.join(".", tablePath))
+            .filter(BigQueryReference::isQualified) // Remove non-qualified tables
+            .collect(Collectors.toSet());
+    this.addTables(List.copyOf(tables));
+  }
+
   /**
    * {@inheritDoc}
    *

diff --git a/...toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/spanner/SpannerCatalog.java b/...toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/spanner/SpannerCatalog.java
@@ -18,6 +18,8 @@
 
 import com.google.cloud.spanner.DatabaseClient;
 import com.google.cloud.spanner.Spanner;
+import com.google.zetasql.Analyzer;
+import com.google.zetasql.AnalyzerOptions;
 import com.google.zetasql.SimpleCatalog;
 import com.google.zetasql.SimpleTable;
 import com.google.zetasql.ZetaSQLBuiltinFunctionOptions;
@@ -32,6 +34,8 @@
 import com.google.zetasql.toolkit.catalog.spanner.exceptions.InvalidSpannerTableName;
 import com.google.zetasql.toolkit.options.SpannerLanguageOptions;
 import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
 
 /** {@link CatalogWrapper} implementation that follows Cloud Spanner semantics */
 public class SpannerCatalog implements CatalogWrapper {
@@ -225,6 +229,24 @@ public void addAllTablesInDatabase() {
                     table, CreateMode.CREATE_OR_REPLACE, CreateScope.CREATE_DEFAULT_SCOPE));
   }
 
+  /**
+   * Adds all the tables used in the provided query to this catalog.
+   *
+   * <p>Uses Analyzer.extractTableNamesFromScript to extract the table names and later uses
+   * this.addTables to add them.
+   *
+   * @param query The SQL query from which to get the tables that should be added to the catalog
+   * @param options The ZetaSQL AnalyzerOptions to use when extracting the table names from the
+   *     query
+   */
+  public void addAllTablesUsedInQuery(String query, AnalyzerOptions options) {
+    Set<String> tables =
+        Analyzer.extractTableNamesFromScript(query, options).stream()
+            .map(tablePath -> String.join(".", tablePath))
+            .collect(Collectors.toSet());
+    this.addTables(List.copyOf(tables));
+  }
+
   @Override
   public void addFunctions(List<String> functions) {
     throw new UnsupportedOperationException(

diff --git a/tools/zetasql-helper/zetasql-toolkit-examples/pom.xml b/tools/zetasql-helper/zetasql-toolkit-examples/pom.xml
@@ -12,7 +12,7 @@
         <maven.compiler.source>11</maven.compiler.source>
         <maven.compiler.target>11</maven.compiler.target>
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
-        <zetasql.toolkit.version>0.2.0</zetasql.toolkit.version>
+        <zetasql.toolkit.version>0.2.1</zetasql.toolkit.version>
         <google.cloud.jib.version>3.3.1</google.cloud.jib.version>
         <container.mainClass/>
     </properties>