From c8c24d9d4fad0a5d684f99386426d39091987524 Mon Sep 17 00:00:00 2001 From: Pablo Paglilla Date: Mon, 6 May 2024 13:02:05 -0300 Subject: [PATCH] Release v0.5.0 (#37) * Update ZetaSQL to version 2023.10.1 * Rewrite queries to fully quote all name paths before analysis * Avoid no longer necessary nesting of tables in catalogs * Update query rewritting to only re-quote name paths that refer to resources (i.e. tables, functions, etc) * Avoid no longer necessary nesting of functions in catalogs * Avoid no longer necessary nesting of TVFs in catalogs * Remove unnecessary slf4j dependency Closes #24 * Reduce code duplication in BigQueryCatalog * Reduce the amount of nesting for procedures in catalogs * Avoid duplicate code when creating different types of resources in catalogs * Use the DOUBLE type kind for BigQuery FLOAT64 columns Closes #25 * Added extractColumnLevelLineage for ResolvedQueryStmt Closes #28 * Changed the access modifier of ParentColumnFinder to public Closes #28 * Make ColumnLineageExtractor accept only concrete statement types Previously, the API for ColumnLineageExtractor used the method ::extractColumnLevelLineage(ResolvedStatement). Since introducing support for ResolvedQueryStmts, which needs specifying an output table separately; maintaining the generic ResolvedStatement API required making it confusing, since it would optionally need to accept an output table. This makes it so that teams building lineage applications need to explicitly determine the statements they support and call the corresponding ::extractColumnLevelLineage() method. Such as ::extractColumnLevelLineage(ResolvedInsertStmt) or ::extractColumnLevelLineage(ResolvedQueryStmt, String). * Bump development version to 0.5.0-SNAPSHOT * vuln-fix: Use HTTPS instead of HTTP to resolve deps CVE-2021-26291 (#30) This fixes a security vulnerability in this project where the `pom.xml` files were configuring Maven to resolve dependencies over HTTP instead of HTTPS. Weakness: CWE-829: Inclusion of Functionality from Untrusted Control Sphere Severity: High CVSS: 8.1 Detection: CodeQL & OpenRewrite (https://app.moderne.io/recipes/org.openrewrite.maven.security.UseHttpsForRepositories) Reported-by: Jonathan Leitschuh Bug-tracker: https://github.com/JLLeitschuh/security-research/issues/8 Detection: CodeQL (https://codeql.github.com/codeql-query-help/java/java-maven-non-https-url/) & OpenRewrite (https://app.moderne.io/recipes/org.openrewrite.maven.security.UseHttpsForRepositories) Reported-by: Jonathan Leitschuh Bug-tracker: https://github.com/JLLeitschuh/security-research/issues/8 Use this link to re-run the recipe: https://app.moderne.io/recipes/org.openrewrite.maven.security.UseHttpsForRepositories?organizationId=R29vZ2xl Co-authored-by: Moderne * Make the type parser case insensitive (#35) The type parser was previously case sensitive, while SQL types are case insensitive. This went unnoticed for a while since upper-cased types are usually always used, but is fundamentally incorrect. Fixes #32 * Add reflection-based patching of GRPC's default max nesting depth (#36) ZetaSQL's Java API uses a GRPC service to call into the actual C++ implementation of ZetaSQL. By default, the serialization logic of that communication allows for a nesting depth in protobuf messages of up to 100. However, long queries can exceed that level of nesting and as a result cannot be analyzed by default. This implements a reflection-based patch that allows users to override that limit to a greater number. This is brittle by design and should be used with caution. Fixes #31 * Upgrade to zetasql-2024-03-01 and bump deps (#33) * Upgrade to zetasql-2024-03-01 and bump deps * Enable all features * Rollback Mockito to version 4.11.0 * Remove some v1.4 language options not supported by BigQuery --------- Co-authored-by: Pablo Paglilla * Add missing @return in catalog Javadocs * Update version to v0.5.0 --------- Co-authored-by: Dion Ricky Saputra Co-authored-by: Jonathan Leitschuh Co-authored-by: Moderne Co-authored-by: Erlend Hamnaberg --- README.md | 4 +- pom.xml | 21 +- zetasql-toolkit-bigquery/pom.xml | 4 +- .../bigquery/BigQueryAPIResourceProvider.java | 2 +- .../catalog/bigquery/BigQueryBuiltIns.java | 18 +- .../catalog/bigquery/BigQueryCatalog.java | 181 +++--- .../catalog/bigquery/BigQueryReference.java | 4 + .../options/BigQueryLanguageOptions.java | 288 +++++----- .../catalog/bigquery/BigQueryCatalogTest.java | 310 +++++----- zetasql-toolkit-core/pom.xml | 12 +- .../catalog/typeparser/ZetaSQLTypeGrammar.g4 | 65 ++- .../com/google/zetasql/SimpleCatalogUtil.java | 32 ++ .../zetasql/toolkit/StatementRewriter.java | 538 ++++++++++++++++++ .../toolkit/ZetaSQLToolkitAnalyzer.java | 9 +- .../toolkit/catalog/CatalogOperations.java | 461 ++++++--------- .../catalog/basic/BasicCatalogWrapper.java | 52 +- .../CatalogResourceDoesNotExist.java | 26 + .../catalog/typeparser/ZetaSQLTypeParser.java | 2 +- .../tools/lineage/ColumnLineageExtractor.java | 55 +- .../tools/lineage/ParentColumnFinder.java | 2 +- .../toolkit/tools/patch/ZetaSQLPatcher.java | 142 +++++ .../catalog/CatalogOperationsTest.java | 247 +++----- .../typeparser/ZetaSQLTypeParserTest.java | 38 +- .../tools/patch/ZetaSQLPatcherTest.java | 68 +++ zetasql-toolkit-examples/pom.xml | 4 +- .../examples/ExtractColumnLevelLineage.java | 20 +- zetasql-toolkit-spanner/pom.xml | 4 +- .../catalog/spanner/SpannerCatalog.java | 18 +- 28 files changed, 1606 insertions(+), 1021 deletions(-) create mode 100644 zetasql-toolkit-core/src/main/java/com/google/zetasql/SimpleCatalogUtil.java create mode 100644 zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/StatementRewriter.java create mode 100644 zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/exceptions/CatalogResourceDoesNotExist.java create mode 100644 zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/tools/patch/ZetaSQLPatcher.java create mode 100644 zetasql-toolkit-core/src/test/java/com/google/zetasql/toolkit/tools/patch/ZetaSQLPatcherTest.java diff --git a/README.md b/README.md index 1d82b33..6255262 100644 --- a/README.md +++ b/README.md @@ -30,7 +30,7 @@ When analyzing queries using BigQuery semantics, you need to: com.google.zetasql.toolkit zetasql-toolkit-bigquery - 0.4.1 + 0.5.0 ``` @@ -113,7 +113,7 @@ Similarly, when analyzing queries using Spanner semantics, you need to: com.google.zetasql.toolkit zetasql-toolkit-spanner - 0.4.1 + 0.5.0 ``` diff --git a/pom.xml b/pom.xml index 595c5a6..b97818f 100644 --- a/pom.xml +++ b/pom.xml @@ -21,7 +21,7 @@ com.google.zetasql.toolkit zetasql-toolkit - 0.4.1 + 0.5.0 pom ${project.groupId}:${project.artifactId} @@ -63,11 +63,10 @@ false true - 2023.04.1 - 26.15.0 - 1.7.25 + 2024.03.1 + 26.37.0 - 5.9.3 + 5.10.2 4.11.0 3.3.0 @@ -114,16 +113,6 @@ zetasql-jni-channel ${zetasql.version} - - org.slf4j - slf4j-api - ${slf4j.version} - - - org.slf4j - slf4j-simple - ${slf4j.version} - org.mockito mockito-core @@ -205,7 +194,7 @@ sonatype-nexus-snapshots Sonatype Nexus Snapshots - http://oss.sonatype.org/content/repositories/snapshots + https://oss.sonatype.org/content/repositories/snapshots diff --git a/zetasql-toolkit-bigquery/pom.xml b/zetasql-toolkit-bigquery/pom.xml index f444ba7..dc2d9b0 100644 --- a/zetasql-toolkit-bigquery/pom.xml +++ b/zetasql-toolkit-bigquery/pom.xml @@ -24,12 +24,12 @@ com.google.zetasql.toolkit zetasql-toolkit - 0.4.1 + 0.5.0 ../pom.xml zetasql-toolkit-bigquery - 0.4.1 + 0.5.0 ${project.groupId}:${project.artifactId} diff --git a/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryAPIResourceProvider.java b/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryAPIResourceProvider.java index 0c492cc..6186506 100644 --- a/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryAPIResourceProvider.java +++ b/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryAPIResourceProvider.java @@ -116,7 +116,7 @@ private TypeKind convertBigqueryTypeNameToTypeKind(StandardSQLTypeName bigqueryT case INT64: return TypeKind.TYPE_INT64; case FLOAT64: - return TypeKind.TYPE_FLOAT; + return TypeKind.TYPE_DOUBLE; case NUMERIC: return TypeKind.TYPE_NUMERIC; case BIGNUMERIC: diff --git a/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryBuiltIns.java b/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryBuiltIns.java index 2eb5501..4643c28 100644 --- a/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryBuiltIns.java +++ b/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryBuiltIns.java @@ -135,7 +135,7 @@ class BigQueryBuiltIns { ImmutableList.of( // BQ.ABORT_SESSION([STRING]) new ProcedureInfo( - ImmutableList.of("BQ", "ABORT_SESSION"), + ImmutableList.of("BQ.ABORT_SESSION"), new FunctionSignature( new FunctionArgumentType(TypeFactory.createSimpleType(TypeKind.TYPE_STRING)), ImmutableList.of( @@ -149,7 +149,7 @@ class BigQueryBuiltIns { -1)), // BQ.JOBS.CANCEL(STRING) new ProcedureInfo( - ImmutableList.of("BQ", "JOBS", "CANCEL"), + ImmutableList.of("BQ.JOBS.CANCEL"), new FunctionSignature( new FunctionArgumentType(TypeFactory.createSimpleType(TypeKind.TYPE_STRING)), ImmutableList.of( @@ -163,7 +163,7 @@ class BigQueryBuiltIns { -1)), // BQ.REFRESH_EXTERNAL_METADATA_CACHE(STRING) new ProcedureInfo( - ImmutableList.of("BQ", "REFRESH_EXTERNAL_METADATA_CACHE"), + ImmutableList.of("BQ.REFRESH_EXTERNAL_METADATA_CACHE"), new FunctionSignature( new FunctionArgumentType(TypeFactory.createSimpleType(TypeKind.TYPE_STRING)), ImmutableList.of( @@ -177,7 +177,7 @@ class BigQueryBuiltIns { -1)), // BQ.REFRESH_MATERIALIZED_VIEW(STRING) new ProcedureInfo( - ImmutableList.of("BQ", "REFRESH_MATERIALIZED_VIEW"), + ImmutableList.of("BQ.REFRESH_MATERIALIZED_VIEW"), new FunctionSignature( new FunctionArgumentType(TypeFactory.createSimpleType(TypeKind.TYPE_STRING)), ImmutableList.of( @@ -199,15 +199,9 @@ class BigQueryBuiltIns { public static void addToCatalog(SimpleCatalog catalog) { TYPE_ALIASES.forEach(catalog::addType); FUNCTIONS.forEach(catalog::addFunction); - - for (ProcedureInfo procedureInfo : PROCEDURES) { - List namePath = procedureInfo.getNamePath(); - String procedureName = namePath.get(namePath.size() - 1); - List> procedurePaths = - ImmutableList.of(namePath, ImmutableList.of(procedureName), ImmutableList.of(String.join(".", namePath))); + PROCEDURES.forEach(procedure -> CatalogOperations.createProcedureInCatalog( - catalog, procedurePaths, procedureInfo, CreateMode.CREATE_OR_REPLACE); - } + catalog, procedure.getFullName(), procedure, CreateMode.CREATE_DEFAULT)); } private BigQueryBuiltIns() {} diff --git a/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryCatalog.java b/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryCatalog.java index 839fa2a..db10976 100644 --- a/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryCatalog.java +++ b/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryCatalog.java @@ -119,6 +119,7 @@ private BigQueryCatalog( * * @param defaultProjectId The BigQuery default project id, queries are assumed to be running on * this project + * @return the new BigQueryCatalog instance */ public static BigQueryCatalog usingBigQueryAPI(String defaultProjectId) { BigQueryResourceProvider resourceProvider = BigQueryAPIResourceProvider.buildDefault(); @@ -135,6 +136,7 @@ public static BigQueryCatalog usingBigQueryAPI(String defaultProjectId) { * @param defaultProjectId The BigQuery default project id, queries are assumed to be running on * this project * @param bigQueryClient The BigQuery client to use for accessing the API + * @return the new BigQueryCatalog instance */ public static BigQueryCatalog usingBigQueryAPI(String defaultProjectId, BigQuery bigQueryClient) { BigQueryResourceProvider resourceProvider = BigQueryAPIResourceProvider.build(bigQueryClient); @@ -148,6 +150,7 @@ public static BigQueryCatalog usingBigQueryAPI(String defaultProjectId, BigQuery * @param defaultProjectId The BigQuery default project id, queries are assumed to be running on * this project * @param resources The {@link CatalogResources} object from which this catalog will get resources + * @return the new BigQueryCatalog instance */ public static BigQueryCatalog usingResources(String defaultProjectId, CatalogResources resources) throws JsonParseException { @@ -269,71 +272,36 @@ private void validateNamePathForCreation( } /** - * Creates the list of resource paths that should be used for creating a resource to make sure it - * is always found when analyzing queries. - * - *

Given the way the ZetaSQL {@link SimpleCatalog} resolves names, different ways of - * referencing resources while querying results in different lookups. For example; "SELECT * FROM - * `A.B`" will look for a table named "A.B", while "SELECT * FROM `A`.`B`" will look for a table - * named "B" on a catalog named "A". - * - *

Because of the previous point, a table or function being created needs to be registered in - * multiple paths. This method creates all those distinct paths. Given the resource - * "project.dataset.resource", this method will create these paths: - * - *

    - *
  • ["project.dataset.resource"] - *
  • ["project", "dataset", "resource"] - *
  • ["project", "dataset.resource"] - *
  • ["project.dataset", "resource"] - *
  • ["dataset.resource"] if the resource project is this catalog's default project id - *
  • ["dataset", "resource"] if the resource project is this catalog's default project id - *
- * - * @param reference The BigQueryReference for the resource that needs to be created - * @return All the distinct name paths at which the resource should be created - */ - private List> buildCatalogPathsForResource(BigQueryReference reference) { - String projectId = reference.getProjectId(); - String datasetName = reference.getDatasetId(); - String resourceName = reference.getResourceName(); - - List> resourcePaths = - ImmutableList.of( - ImmutableList.of(projectId, datasetName, resourceName), // format: project.dataset.table format - ImmutableList.of( - projectId - + "." - + datasetName - + "." - + resourceName), // format: `project.dataset.table` - ImmutableList.of(projectId + "." + datasetName, resourceName), // format: `project.dataset`.table - ImmutableList.of(projectId, datasetName + "." + resourceName) // format: project.`dataset.table` - ); - - List> resourcePathsWithImplicitProject = ImmutableList.of(); - - if (projectId.equals(this.defaultProjectId)) { - resourcePathsWithImplicitProject = - ImmutableList.of( - ImmutableList.of(datasetName, resourceName), // format: dataset.table (project implied) - ImmutableList.of(datasetName + "." + resourceName) // format: `dataset.table` (project implied) - ); + * Returns the names a resource referenced by the provided string should have in the underlying + * {@link SimpleCatalog}. + * + *

If the reference is qualified, its complete name in the catalog will be in the form of + * "project.dataset.resource". If the resource is in the default project, an additional name + * in the form of "dataset.resource" will be returned. + * + *

For unqualified resource (e.g. temporary tables), the name will be used as-is. + * + * @param reference The string used to reference the resource (e.g. "project.dataset.table", + * "dataset.function", "tableName") + * @return The list of names the resource should have in the underlying {@link SimpleCatalog} + */ + private List buildCatalogNamesForResource(String reference) { + boolean isQualified = BigQueryReference.isQualified(reference); + + if (!isQualified) { + return ImmutableList.of(reference); } - return Stream.concat(resourcePaths.stream(), resourcePathsWithImplicitProject.stream()) - .collect(Collectors.toList()); - } + BigQueryReference parsedReference = + BigQueryReference.from(this.defaultProjectId, reference); + boolean isInDefaultProject = + parsedReference.getProjectId().equalsIgnoreCase(this.defaultProjectId); - /** @see #buildCatalogPathsForResource(BigQueryReference) */ - private List> buildCatalogPathsForResource(String referenceStr) { - BigQueryReference reference = BigQueryReference.from(this.defaultProjectId, referenceStr); - return this.buildCatalogPathsForResource(reference); - } - - /** @see #buildCatalogPathsForResource(BigQueryReference) */ - private List> buildCatalogPathsForResource(List resourcePath) { - return this.buildCatalogPathsForResource(String.join(".", resourcePath)); + if (isInDefaultProject) { + return ImmutableList.of(parsedReference.getFullName(), parsedReference.getNameWithDataset()); + } else { + return ImmutableList.of(parsedReference.getFullName()); + } } private CatalogResourceAlreadyExists addCaseInsensitivityWarning( @@ -351,10 +319,10 @@ private CatalogResourceAlreadyExists addCaseInsensitivityWarning( /** * {@inheritDoc} * - *

Multiple copies of the registered {@link SimpleTable} will be created in the Catalog to - * comply with BigQuery name resolution semantics. + *

If the table is not temporary and is in the catalog's default project, it will be + * registered twice. Once as "project.dataset.table" and once as "dataset.table". That way, + * queries that omit the project when referencing the table can be analyzed. * - * @see #buildCatalogPathsForResource(BigQueryReference) * @throws BigQueryCreateError if a pre-create validation fails * @throws CatalogResourceAlreadyExists if the table already exists and CreateMode != * CREATE_OR_REPLACE @@ -366,16 +334,14 @@ public void register(SimpleTable table, CreateMode createMode, CreateScope creat ImmutableList.of(CreateScope.CREATE_DEFAULT_SCOPE, CreateScope.CREATE_TEMP), table.getFullName(), "table"); - this.validateNamePathForCreation(ImmutableList.of(table.getFullName()), createScope, "table"); + this.validateNamePathForCreation( + ImmutableList.of(table.getFullName()), createScope, "table"); - List> tablePaths = - createScope.equals(CreateScope.CREATE_TEMP) - ? ImmutableList.of(ImmutableList.of(table.getName())) - : this.buildCatalogPathsForResource(table.getFullName()); + List catalogNamesForTable = buildCatalogNamesForResource(table.getFullName()); try { - CatalogOperations.createTableInCatalog( - this.catalog, tablePaths, table.getFullName(), table.getColumnList(), createMode); + catalogNamesForTable.forEach(catalogName -> + CatalogOperations.createTableInCatalog(catalog, catalogName, table, createMode)); } catch (CatalogResourceAlreadyExists alreadyExists) { throw this.addCaseInsensitivityWarning(alreadyExists); } @@ -384,14 +350,14 @@ public void register(SimpleTable table, CreateMode createMode, CreateScope creat /** * {@inheritDoc} * - *

Multiple copies of the registered {@link FunctionInfo} will be created in the Catalog to - * comply with BigQuery name resolution semantics. + *

If the function is not temporary and is in the catalog's default project, it will be + * registered twice. Once as "project.dataset.function" and once as "dataset.function". That way, + * queries that omit the project when referencing the table can be analyzed. * *

If any function signature which does not include a valid return type, this method will * attempt to infer it. If inference is not possible, {@link MissingFunctionResultType} will be * thrown. * - * @see #buildCatalogPathsForResource(BigQueryReference) * @throws BigQueryCreateError if a pre-create validation fails * @throws CatalogResourceAlreadyExists if the function already exists and CreateMode != * CREATE_OR_REPLACE @@ -414,14 +380,12 @@ public void register(FunctionInfo function, CreateMode createMode, CreateScope c FunctionResultTypeResolver.resolveFunctionReturnTypes( function, BigQueryLanguageOptions.get(), this.catalog); - List> functionPaths = - createScope.equals(CreateScope.CREATE_TEMP) - ? ImmutableList.of(functionNamePath) - : this.buildCatalogPathsForResource(functionNamePath); + List catalogNamesForFunction = buildCatalogNamesForResource(fullName); try { - CatalogOperations.createFunctionInCatalog( - this.catalog, functionPaths, resolvedFunction, createMode); + catalogNamesForFunction.forEach(catalogName -> + CatalogOperations.createFunctionInCatalog( + catalog, catalogName, resolvedFunction, createMode)); } catch (CatalogResourceAlreadyExists alreadyExists) { throw this.addCaseInsensitivityWarning(alreadyExists); } @@ -430,6 +394,10 @@ public void register(FunctionInfo function, CreateMode createMode, CreateScope c /** * {@inheritDoc} * + *

If the function in the catalog's default project, it will be registered twice. Once as + * "project.dataset.function" and once as "dataset.function". That way, queries that omit the + * project when referencing the table can be analyzed. + * * @throws BigQueryCreateError if a pre-create validation fails * @throws CatalogResourceAlreadyExists if the function already exists and CreateMode != * CREATE_OR_REPLACE @@ -446,11 +414,11 @@ public void register(TVFInfo tvfInfo, CreateMode createMode, CreateScope createS FunctionResultTypeResolver.resolveTVFOutputSchema( tvfInfo, BigQueryLanguageOptions.get(), this.catalog); - List> functionPaths = this.buildCatalogPathsForResource(fullName); + List catalogNamesForFunction = buildCatalogNamesForResource(fullName); try { - CatalogOperations.createTVFInCatalog( - this.catalog, functionPaths, resolvedTvfInfo, createMode); + catalogNamesForFunction.forEach(catalogName -> + CatalogOperations.createTVFInCatalog(catalog, catalogName, resolvedTvfInfo, createMode)); } catch (CatalogResourceAlreadyExists alreadyExists) { throw this.addCaseInsensitivityWarning(alreadyExists); } @@ -459,10 +427,10 @@ public void register(TVFInfo tvfInfo, CreateMode createMode, CreateScope createS /** * {@inheritDoc} * - *

Multiple copies of the registered {@link Procedure} will be created in the Catalog to comply - * with BigQuery name resolution semantics. + *

If the procedure in the catalog's default project, it will be registered twice. Once as + * "project.dataset.procedure" and once as "dataset.procedure". That way, queries that omit the + * project when referencing the table can be analyzed. * - * @see #buildCatalogPathsForResource(BigQueryReference) * @throws BigQueryCreateError if a pre-create validation fails * @throws CatalogResourceAlreadyExists if the precedure already exists and CreateMode != * CREATE_OR_REPLACE @@ -476,11 +444,12 @@ public void register( createScope, ImmutableList.of(CreateScope.CREATE_DEFAULT_SCOPE), fullName, "procedure"); this.validateNamePathForCreation(ImmutableList.of(fullName), createScope, "procedure"); - List> procedurePaths = this.buildCatalogPathsForResource(fullName); + List catalogNamesForProcedure = buildCatalogNamesForResource(fullName); try { - CatalogOperations.createProcedureInCatalog( - this.catalog, procedurePaths, procedureInfo, createMode); + catalogNamesForProcedure.forEach(catalogName -> + CatalogOperations.createProcedureInCatalog( + catalog, catalogName, procedureInfo, createMode)); } catch (CatalogResourceAlreadyExists alreadyExists) { throw this.addCaseInsensitivityWarning(alreadyExists); } @@ -515,38 +484,34 @@ public void register(Constant constant) { @Override public void removeTable(String tableReference) { - boolean isQualified = tableReference.split("\\.").length > 1; - - List> tablePaths = - !isQualified - ? ImmutableList.of(ImmutableList.of(tableReference)) - : this.buildCatalogPathsForResource(tableReference); + List catalogNamesForTable = buildCatalogNamesForResource(tableReference); - CatalogOperations.deleteTableFromCatalog(this.catalog, tablePaths); + catalogNamesForTable.forEach(catalogName -> + CatalogOperations.deleteTableFromCatalog(catalog, catalogName)); } @Override public void removeFunction(String functionReference) { - boolean isQualified = functionReference.split("\\.").length > 1; + List catalogNamesForFunction = buildCatalogNamesForResource(functionReference); - List> functionPaths = - !isQualified - ? ImmutableList.of(ImmutableList.of(functionReference)) - : this.buildCatalogPathsForResource(functionReference); - - CatalogOperations.deleteFunctionFromCatalog(this.catalog, functionPaths); + catalogNamesForFunction.forEach(catalogName -> + CatalogOperations.deleteFunctionFromCatalog(catalog, catalogName)); } @Override public void removeTVF(String functionReference) { - List> functionPaths = this.buildCatalogPathsForResource(functionReference); - CatalogOperations.deleteTVFFromCatalog(this.catalog, functionPaths); + List catalogNamesForFunction = buildCatalogNamesForResource(functionReference); + + catalogNamesForFunction.forEach(catalogName -> + CatalogOperations.deleteTVFFromCatalog(catalog, catalogName)); } @Override public void removeProcedure(String procedureReference) { - List> functionPaths = this.buildCatalogPathsForResource(procedureReference); - CatalogOperations.deleteProcedureFromCatalog(this.catalog, functionPaths); + List catalogNamesForProcedure = buildCatalogNamesForResource(procedureReference); + + catalogNamesForProcedure.forEach(catalogName -> + CatalogOperations.deleteProcedureFromCatalog(catalog, catalogName)); } /** diff --git a/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryReference.java b/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryReference.java index 5ebbc41..90c3ea3 100644 --- a/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryReference.java +++ b/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryReference.java @@ -107,6 +107,10 @@ public String getFullName() { "%s.%s.%s", this.getProjectId(), this.getDatasetId(), this.getResourceName()); } + public String getNameWithDataset() { + return String.format("%s.%s", this.getDatasetId(), this.getResourceName()); + } + public ImmutableList getNamePath() { return ImmutableList.of(this.getProjectId(), this.getDatasetId(), this.getResourceName()); } diff --git a/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/options/BigQueryLanguageOptions.java b/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/options/BigQueryLanguageOptions.java index cae90bd..061bb40 100644 --- a/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/options/BigQueryLanguageOptions.java +++ b/zetasql-toolkit-bigquery/src/main/java/com/google/zetasql/toolkit/options/BigQueryLanguageOptions.java @@ -25,151 +25,155 @@ public class BigQueryLanguageOptions { - private static final LanguageOptions languageOptions = new LanguageOptions(); + private static final LanguageOptions languageOptions = new LanguageOptions(); - static { - languageOptions.setNameResolutionMode(NameResolutionMode.NAME_RESOLUTION_DEFAULT); - languageOptions.setProductMode(ProductMode.PRODUCT_EXTERNAL); + static { + languageOptions.setNameResolutionMode(NameResolutionMode.NAME_RESOLUTION_DEFAULT); + languageOptions.setProductMode(ProductMode.PRODUCT_EXTERNAL); - languageOptions.setEnabledLanguageFeatures( - ImmutableSet.of( - LanguageFeature.FEATURE_ALLOW_MISSING_PATH_EXPRESSION_IN_ALTER_DDL, - LanguageFeature.FEATURE_ALTER_COLUMN_SET_DATA_TYPE, - LanguageFeature.FEATURE_ALTER_TABLE_RENAME_COLUMN, - LanguageFeature.FEATURE_ANALYTIC_FUNCTIONS, - LanguageFeature.FEATURE_ANONYMIZATION, - LanguageFeature.FEATURE_BIGNUMERIC_TYPE, - LanguageFeature.FEATURE_CBRT_FUNCTIONS, - LanguageFeature.FEATURE_CREATE_EXTERNAL_TABLE_WITH_CONNECTION, - LanguageFeature.FEATURE_CREATE_EXTERNAL_TABLE_WITH_PARTITION_COLUMNS, - LanguageFeature.FEATURE_CREATE_EXTERNAL_TABLE_WITH_TABLE_ELEMENT_LIST, - LanguageFeature.FEATURE_CREATE_MATERIALIZED_VIEW_CLUSTER_BY, - LanguageFeature.FEATURE_CREATE_MATERIALIZED_VIEW_PARTITION_BY, - LanguageFeature.FEATURE_CREATE_SNAPSHOT_TABLE, - LanguageFeature.FEATURE_CREATE_TABLE_AS_SELECT_COLUMN_LIST, - LanguageFeature.FEATURE_CREATE_TABLE_CLONE, - LanguageFeature.FEATURE_CREATE_TABLE_CLUSTER_BY, - LanguageFeature.FEATURE_CREATE_TABLE_COPY, - LanguageFeature.FEATURE_CREATE_TABLE_FIELD_ANNOTATIONS, - LanguageFeature.FEATURE_CREATE_TABLE_FUNCTION, - LanguageFeature.FEATURE_CREATE_TABLE_LIKE, - LanguageFeature.FEATURE_CREATE_TABLE_NOT_NULL, - LanguageFeature.FEATURE_CREATE_TABLE_PARTITION_BY, - LanguageFeature.FEATURE_CREATE_VIEW_WITH_COLUMN_LIST, - LanguageFeature.FEATURE_DML_UPDATE_WITH_JOIN, - LanguageFeature.FEATURE_ENCRYPTION, - LanguageFeature.FEATURE_FOREIGN_KEYS, - LanguageFeature.FEATURE_GEOGRAPHY, - LanguageFeature.FEATURE_GROUP_BY_ROLLUP, - LanguageFeature.FEATURE_INTERVAL_TYPE, - LanguageFeature.FEATURE_INVERSE_TRIG_FUNCTIONS, - LanguageFeature.FEATURE_JSON_ARRAY_FUNCTIONS, - LanguageFeature.FEATURE_JSON_TYPE, - LanguageFeature.FEATURE_JSON_VALUE_EXTRACTION_FUNCTIONS, - LanguageFeature.FEATURE_NAMED_ARGUMENTS, - LanguageFeature.FEATURE_NUMERIC_TYPE, - LanguageFeature.FEATURE_PARAMETERIZED_TYPES, - LanguageFeature.FEATURE_PARAMETERS_IN_GRANTEE_LIST, - LanguageFeature.FEATURE_ROUND_WITH_ROUNDING_MODE, - LanguageFeature.FEATURE_TABLESAMPLE, - LanguageFeature.FEATURE_TABLE_VALUED_FUNCTIONS, - LanguageFeature.FEATURE_TEMPLATE_FUNCTIONS, - LanguageFeature.FEATURE_TIMESTAMP_NANOS, - LanguageFeature.FEATURE_TIME_BUCKET_FUNCTIONS, - LanguageFeature.FEATURE_UNENFORCED_PRIMARY_KEYS, - LanguageFeature.FEATURE_V_1_1_HAVING_IN_AGGREGATE, - LanguageFeature.FEATURE_V_1_1_LIMIT_IN_AGGREGATE, - LanguageFeature.FEATURE_V_1_1_NULL_HANDLING_MODIFIER_IN_AGGREGATE, - LanguageFeature.FEATURE_V_1_1_NULL_HANDLING_MODIFIER_IN_ANALYTIC, - LanguageFeature.FEATURE_V_1_1_ORDER_BY_IN_AGGREGATE, - LanguageFeature.FEATURE_V_1_1_SELECT_STAR_EXCEPT_REPLACE, - LanguageFeature.FEATURE_V_1_1_WITH_ON_SUBQUERY, - LanguageFeature.FEATURE_V_1_2_CIVIL_TIME, - LanguageFeature.FEATURE_V_1_2_SAFE_FUNCTION_CALL, - LanguageFeature.FEATURE_V_1_2_WEEK_WITH_WEEKDAY, - LanguageFeature.FEATURE_V_1_3_ADDITIONAL_STRING_FUNCTIONS, - LanguageFeature.FEATURE_V_1_3_ALLOW_DASHES_IN_TABLE_NAME, - LanguageFeature.FEATURE_V_1_3_ALLOW_REGEXP_EXTRACT_OPTIONALS, - LanguageFeature.FEATURE_V_1_3_ANNOTATION_FRAMEWORK, - LanguageFeature.FEATURE_V_1_3_CASE_STMT, - LanguageFeature.FEATURE_V_1_3_COLLATION_SUPPORT, - LanguageFeature.FEATURE_V_1_3_COLUMN_DEFAULT_VALUE, - LanguageFeature.FEATURE_V_1_3_CONCAT_MIXED_TYPES, - LanguageFeature.FEATURE_V_1_3_DATE_ARITHMETICS, - LanguageFeature.FEATURE_V_1_3_DATE_TIME_CONSTRUCTORS, - LanguageFeature.FEATURE_V_1_3_DECIMAL_ALIAS, - LanguageFeature.FEATURE_V_1_3_EXTENDED_DATE_TIME_SIGNATURES, - LanguageFeature.FEATURE_V_1_3_EXTENDED_GEOGRAPHY_PARSERS, - LanguageFeature.FEATURE_V_1_3_FORMAT_IN_CAST, - LanguageFeature.FEATURE_V_1_3_FOR_IN, - LanguageFeature.FEATURE_V_1_3_IS_DISTINCT, - LanguageFeature.FEATURE_V_1_3_LIKE_ANY_SOME_ALL, - LanguageFeature.FEATURE_V_1_3_NULLS_FIRST_LAST_IN_ORDER_BY, - LanguageFeature.FEATURE_V_1_3_OMIT_INSERT_COLUMN_LIST, - LanguageFeature.FEATURE_V_1_3_PIVOT, - LanguageFeature.FEATURE_V_1_3_QUALIFY, - LanguageFeature.FEATURE_V_1_3_REMOTE_FUNCTION, - LanguageFeature.FEATURE_V_1_3_REPEAT, - LanguageFeature.FEATURE_V_1_3_SCRIPT_LABEL, - LanguageFeature.FEATURE_V_1_3_UNNEST_AND_FLATTEN_ARRAYS, - LanguageFeature.FEATURE_V_1_3_UNPIVOT, - LanguageFeature.FEATURE_V_1_3_WITH_GROUP_ROWS, - LanguageFeature.FEATURE_V_1_3_WITH_RECURSIVE)); + languageOptions.setEnabledLanguageFeatures( + ImmutableSet.of( + LanguageFeature.FEATURE_ALLOW_MISSING_PATH_EXPRESSION_IN_ALTER_DDL, + LanguageFeature.FEATURE_ALTER_COLUMN_SET_DATA_TYPE, + LanguageFeature.FEATURE_ALTER_TABLE_RENAME_COLUMN, + LanguageFeature.FEATURE_ANALYTIC_FUNCTIONS, + LanguageFeature.FEATURE_ANONYMIZATION, + LanguageFeature.FEATURE_BIGNUMERIC_TYPE, + LanguageFeature.FEATURE_CBRT_FUNCTIONS, + LanguageFeature.FEATURE_CREATE_EXTERNAL_TABLE_WITH_CONNECTION, + LanguageFeature.FEATURE_CREATE_EXTERNAL_TABLE_WITH_PARTITION_COLUMNS, + LanguageFeature.FEATURE_CREATE_EXTERNAL_TABLE_WITH_TABLE_ELEMENT_LIST, + LanguageFeature.FEATURE_CREATE_MATERIALIZED_VIEW_CLUSTER_BY, + LanguageFeature.FEATURE_CREATE_MATERIALIZED_VIEW_PARTITION_BY, + LanguageFeature.FEATURE_CREATE_SNAPSHOT_TABLE, + LanguageFeature.FEATURE_CREATE_TABLE_AS_SELECT_COLUMN_LIST, + LanguageFeature.FEATURE_CREATE_TABLE_CLONE, + LanguageFeature.FEATURE_CREATE_TABLE_CLUSTER_BY, + LanguageFeature.FEATURE_CREATE_TABLE_COPY, + LanguageFeature.FEATURE_CREATE_TABLE_FIELD_ANNOTATIONS, + LanguageFeature.FEATURE_CREATE_TABLE_FUNCTION, + LanguageFeature.FEATURE_CREATE_TABLE_LIKE, + LanguageFeature.FEATURE_CREATE_TABLE_NOT_NULL, + LanguageFeature.FEATURE_CREATE_TABLE_PARTITION_BY, + LanguageFeature.FEATURE_CREATE_VIEW_WITH_COLUMN_LIST, + LanguageFeature.FEATURE_DML_UPDATE_WITH_JOIN, + LanguageFeature.FEATURE_ENCRYPTION, + LanguageFeature.FEATURE_FOREIGN_KEYS, + LanguageFeature.FEATURE_GEOGRAPHY, + LanguageFeature.FEATURE_GROUP_BY_ROLLUP, + LanguageFeature.FEATURE_INTERVAL_TYPE, + LanguageFeature.FEATURE_INVERSE_TRIG_FUNCTIONS, + LanguageFeature.FEATURE_JSON_ARRAY_FUNCTIONS, + LanguageFeature.FEATURE_JSON_TYPE, + LanguageFeature.FEATURE_JSON_VALUE_EXTRACTION_FUNCTIONS, + LanguageFeature.FEATURE_NAMED_ARGUMENTS, + LanguageFeature.FEATURE_NUMERIC_TYPE, + LanguageFeature.FEATURE_PARAMETERIZED_TYPES, + LanguageFeature.FEATURE_PARAMETERS_IN_GRANTEE_LIST, + LanguageFeature.FEATURE_ROUND_WITH_ROUNDING_MODE, + LanguageFeature.FEATURE_TABLESAMPLE, + LanguageFeature.FEATURE_TABLE_VALUED_FUNCTIONS, + LanguageFeature.FEATURE_TEMPLATE_FUNCTIONS, + LanguageFeature.FEATURE_TIMESTAMP_NANOS, + LanguageFeature.FEATURE_TIME_BUCKET_FUNCTIONS, + LanguageFeature.FEATURE_UNENFORCED_PRIMARY_KEYS, + LanguageFeature.FEATURE_V_1_1_HAVING_IN_AGGREGATE, + LanguageFeature.FEATURE_V_1_1_LIMIT_IN_AGGREGATE, + LanguageFeature.FEATURE_V_1_1_NULL_HANDLING_MODIFIER_IN_AGGREGATE, + LanguageFeature.FEATURE_V_1_1_NULL_HANDLING_MODIFIER_IN_ANALYTIC, + LanguageFeature.FEATURE_V_1_1_ORDER_BY_IN_AGGREGATE, + LanguageFeature.FEATURE_V_1_1_SELECT_STAR_EXCEPT_REPLACE, + LanguageFeature.FEATURE_V_1_1_WITH_ON_SUBQUERY, + LanguageFeature.FEATURE_V_1_2_CIVIL_TIME, + LanguageFeature.FEATURE_V_1_2_SAFE_FUNCTION_CALL, + LanguageFeature.FEATURE_V_1_2_WEEK_WITH_WEEKDAY, + LanguageFeature.FEATURE_V_1_3_ADDITIONAL_STRING_FUNCTIONS, + LanguageFeature.FEATURE_V_1_3_ALLOW_DASHES_IN_TABLE_NAME, + LanguageFeature.FEATURE_V_1_3_ALLOW_REGEXP_EXTRACT_OPTIONALS, + LanguageFeature.FEATURE_V_1_3_ANNOTATION_FRAMEWORK, + LanguageFeature.FEATURE_V_1_3_CASE_STMT, + LanguageFeature.FEATURE_V_1_3_COLLATION_SUPPORT, + LanguageFeature.FEATURE_V_1_3_COLUMN_DEFAULT_VALUE, + LanguageFeature.FEATURE_V_1_3_CONCAT_MIXED_TYPES, + LanguageFeature.FEATURE_V_1_3_DATE_ARITHMETICS, + LanguageFeature.FEATURE_V_1_3_DATE_TIME_CONSTRUCTORS, + LanguageFeature.FEATURE_V_1_3_DECIMAL_ALIAS, + LanguageFeature.FEATURE_V_1_3_EXTENDED_DATE_TIME_SIGNATURES, + LanguageFeature.FEATURE_V_1_3_EXTENDED_GEOGRAPHY_PARSERS, + LanguageFeature.FEATURE_V_1_3_FORMAT_IN_CAST, + LanguageFeature.FEATURE_V_1_3_FOR_IN, + LanguageFeature.FEATURE_V_1_3_IS_DISTINCT, + LanguageFeature.FEATURE_V_1_3_LIKE_ANY_SOME_ALL, + LanguageFeature.FEATURE_V_1_3_NULLS_FIRST_LAST_IN_ORDER_BY, + LanguageFeature.FEATURE_V_1_3_OMIT_INSERT_COLUMN_LIST, + LanguageFeature.FEATURE_V_1_3_PIVOT, + LanguageFeature.FEATURE_V_1_3_QUALIFY, + LanguageFeature.FEATURE_V_1_3_REMOTE_FUNCTION, + LanguageFeature.FEATURE_V_1_3_REPEAT, + LanguageFeature.FEATURE_V_1_3_SCRIPT_LABEL, + LanguageFeature.FEATURE_V_1_3_UNNEST_AND_FLATTEN_ARRAYS, + LanguageFeature.FEATURE_V_1_3_UNPIVOT, + LanguageFeature.FEATURE_V_1_3_WITH_GROUP_ROWS, + LanguageFeature.FEATURE_V_1_3_WITH_RECURSIVE, + LanguageFeature.FEATURE_V_1_4_ARRAY_AGGREGATION_FUNCTIONS, + LanguageFeature.FEATURE_V_1_4_REMOTE_MODEL, + LanguageFeature.FEATURE_V_1_4_GROUPING_BUILTIN, + LanguageFeature.FEATURE_V_1_4_GROUPING_SETS, + LanguageFeature.FEATURE_V_1_4_GROUP_BY_ALL)); - languageOptions.setSupportedStatementKinds( - ImmutableSet.of( - ResolvedNodeKind.RESOLVED_ADD_CONSTRAINT_ACTION, - ResolvedNodeKind.RESOLVED_ALTER_ENTITY_STMT, - ResolvedNodeKind.RESOLVED_ALTER_MATERIALIZED_VIEW_STMT, - ResolvedNodeKind.RESOLVED_ALTER_MODEL_STMT, - ResolvedNodeKind.RESOLVED_ALTER_SCHEMA_STMT, - ResolvedNodeKind.RESOLVED_ALTER_TABLE_STMT, - ResolvedNodeKind.RESOLVED_ALTER_VIEW_STMT, - ResolvedNodeKind.RESOLVED_ASSERT_STMT, - ResolvedNodeKind.RESOLVED_ASSIGNMENT_STMT, - ResolvedNodeKind.RESOLVED_AUX_LOAD_DATA_STMT, - ResolvedNodeKind.RESOLVED_BEGIN_STMT, - ResolvedNodeKind.RESOLVED_CALL_STMT, - ResolvedNodeKind.RESOLVED_COMMIT_STMT, - ResolvedNodeKind.RESOLVED_CREATE_ENTITY_STMT, - ResolvedNodeKind.RESOLVED_CREATE_EXTERNAL_TABLE_STMT, - ResolvedNodeKind.RESOLVED_CREATE_FUNCTION_STMT, - ResolvedNodeKind.RESOLVED_CREATE_INDEX_STMT, - ResolvedNodeKind.RESOLVED_CREATE_MATERIALIZED_VIEW_STMT, - ResolvedNodeKind.RESOLVED_CREATE_MODEL_STMT, - ResolvedNodeKind.RESOLVED_CREATE_PROCEDURE_STMT, - ResolvedNodeKind.RESOLVED_CREATE_ROW_ACCESS_POLICY_STMT, - ResolvedNodeKind.RESOLVED_CREATE_SCHEMA_STMT, - ResolvedNodeKind.RESOLVED_CREATE_SNAPSHOT_TABLE_STMT, - ResolvedNodeKind.RESOLVED_CREATE_TABLE_AS_SELECT_STMT, - ResolvedNodeKind.RESOLVED_CREATE_TABLE_FUNCTION_STMT, - ResolvedNodeKind.RESOLVED_CREATE_TABLE_STMT, - ResolvedNodeKind.RESOLVED_CREATE_VIEW_STMT, - ResolvedNodeKind.RESOLVED_DELETE_STMT, - ResolvedNodeKind.RESOLVED_DROP_CONSTRAINT_ACTION, - ResolvedNodeKind.RESOLVED_DROP_FUNCTION_STMT, - ResolvedNodeKind.RESOLVED_DROP_MATERIALIZED_VIEW_STMT, - ResolvedNodeKind.RESOLVED_DROP_ROW_ACCESS_POLICY_STMT, - ResolvedNodeKind.RESOLVED_DROP_SEARCH_INDEX_STMT, - ResolvedNodeKind.RESOLVED_DROP_SNAPSHOT_TABLE_STMT, - ResolvedNodeKind.RESOLVED_DROP_STMT, - ResolvedNodeKind.RESOLVED_DROP_TABLE_FUNCTION_STMT, - ResolvedNodeKind.RESOLVED_EXECUTE_IMMEDIATE_STMT, - ResolvedNodeKind.RESOLVED_EXPORT_DATA_STMT, - ResolvedNodeKind.RESOLVED_EXPORT_MODEL_STMT, - ResolvedNodeKind.RESOLVED_GRANT_STMT, - ResolvedNodeKind.RESOLVED_INSERT_STMT, - ResolvedNodeKind.RESOLVED_MERGE_STMT, - ResolvedNodeKind.RESOLVED_QUERY_STMT, - ResolvedNodeKind.RESOLVED_REVOKE_STMT, - ResolvedNodeKind.RESOLVED_ROLLBACK_STMT, - ResolvedNodeKind.RESOLVED_TRUNCATE_STMT, - ResolvedNodeKind.RESOLVED_UPDATE_STMT)); + languageOptions.setSupportedStatementKinds( + ImmutableSet.of( + ResolvedNodeKind.RESOLVED_ADD_CONSTRAINT_ACTION, + ResolvedNodeKind.RESOLVED_ALTER_ENTITY_STMT, + ResolvedNodeKind.RESOLVED_ALTER_MATERIALIZED_VIEW_STMT, + ResolvedNodeKind.RESOLVED_ALTER_MODEL_STMT, + ResolvedNodeKind.RESOLVED_ALTER_SCHEMA_STMT, + ResolvedNodeKind.RESOLVED_ALTER_TABLE_STMT, + ResolvedNodeKind.RESOLVED_ALTER_VIEW_STMT, + ResolvedNodeKind.RESOLVED_ASSERT_STMT, + ResolvedNodeKind.RESOLVED_ASSIGNMENT_STMT, + ResolvedNodeKind.RESOLVED_AUX_LOAD_DATA_STMT, + ResolvedNodeKind.RESOLVED_BEGIN_STMT, + ResolvedNodeKind.RESOLVED_CALL_STMT, + ResolvedNodeKind.RESOLVED_COMMIT_STMT, + ResolvedNodeKind.RESOLVED_CREATE_ENTITY_STMT, + ResolvedNodeKind.RESOLVED_CREATE_EXTERNAL_TABLE_STMT, + ResolvedNodeKind.RESOLVED_CREATE_FUNCTION_STMT, + ResolvedNodeKind.RESOLVED_CREATE_INDEX_STMT, + ResolvedNodeKind.RESOLVED_CREATE_MATERIALIZED_VIEW_STMT, + ResolvedNodeKind.RESOLVED_CREATE_MODEL_STMT, + ResolvedNodeKind.RESOLVED_CREATE_PROCEDURE_STMT, + ResolvedNodeKind.RESOLVED_CREATE_ROW_ACCESS_POLICY_STMT, + ResolvedNodeKind.RESOLVED_CREATE_SCHEMA_STMT, + ResolvedNodeKind.RESOLVED_CREATE_SNAPSHOT_TABLE_STMT, + ResolvedNodeKind.RESOLVED_CREATE_TABLE_AS_SELECT_STMT, + ResolvedNodeKind.RESOLVED_CREATE_TABLE_FUNCTION_STMT, + ResolvedNodeKind.RESOLVED_CREATE_TABLE_STMT, + ResolvedNodeKind.RESOLVED_CREATE_VIEW_STMT, + ResolvedNodeKind.RESOLVED_DELETE_STMT, + ResolvedNodeKind.RESOLVED_DROP_CONSTRAINT_ACTION, + ResolvedNodeKind.RESOLVED_DROP_FUNCTION_STMT, + ResolvedNodeKind.RESOLVED_DROP_MATERIALIZED_VIEW_STMT, + ResolvedNodeKind.RESOLVED_DROP_ROW_ACCESS_POLICY_STMT, + ResolvedNodeKind.RESOLVED_DROP_SNAPSHOT_TABLE_STMT, + ResolvedNodeKind.RESOLVED_DROP_STMT, + ResolvedNodeKind.RESOLVED_DROP_TABLE_FUNCTION_STMT, + ResolvedNodeKind.RESOLVED_EXECUTE_IMMEDIATE_STMT, + ResolvedNodeKind.RESOLVED_EXPORT_DATA_STMT, + ResolvedNodeKind.RESOLVED_EXPORT_MODEL_STMT, + ResolvedNodeKind.RESOLVED_GRANT_STMT, + ResolvedNodeKind.RESOLVED_INSERT_STMT, + ResolvedNodeKind.RESOLVED_MERGE_STMT, + ResolvedNodeKind.RESOLVED_QUERY_STMT, + ResolvedNodeKind.RESOLVED_REVOKE_STMT, + ResolvedNodeKind.RESOLVED_ROLLBACK_STMT, + ResolvedNodeKind.RESOLVED_TRUNCATE_STMT, + ResolvedNodeKind.RESOLVED_UPDATE_STMT)); - languageOptions.enableReservableKeyword("QUALIFY"); - } + languageOptions.enableReservableKeyword("QUALIFY"); + } - public static LanguageOptions get() { - return languageOptions; - } + public static LanguageOptions get() { + return languageOptions; + } } diff --git a/zetasql-toolkit-bigquery/src/test/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryCatalogTest.java b/zetasql-toolkit-bigquery/src/test/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryCatalogTest.java index 9d4f20c..7906bf9 100644 --- a/zetasql-toolkit-bigquery/src/test/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryCatalogTest.java +++ b/zetasql-toolkit-bigquery/src/test/java/com/google/zetasql/toolkit/catalog/bigquery/BigQueryCatalogTest.java @@ -21,7 +21,17 @@ import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; -import com.google.zetasql.*; +import com.google.zetasql.Function; +import com.google.zetasql.FunctionArgumentType; +import com.google.zetasql.FunctionSignature; +import com.google.zetasql.NotFoundException; +import com.google.zetasql.SimpleCatalog; +import com.google.zetasql.SimpleColumn; +import com.google.zetasql.SimpleTable; +import com.google.zetasql.TVFRelation; +import com.google.zetasql.Table; +import com.google.zetasql.Type; +import com.google.zetasql.TypeFactory; import com.google.zetasql.ZetaSQLFunctions.FunctionEnums.Mode; import com.google.zetasql.ZetaSQLFunctions.SignatureArgumentKind; import com.google.zetasql.ZetaSQLType.TypeKind; @@ -35,12 +45,9 @@ import com.google.zetasql.toolkit.catalog.exceptions.CatalogResourceAlreadyExists; import java.util.ArrayList; import java.util.List; -import java.util.stream.Stream; -import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; -import org.junit.jupiter.api.function.Executable; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; @@ -50,24 +57,6 @@ public class BigQueryCatalogTest { final String testProjectId = "test-project-id"; BigQueryCatalog bigQueryCatalog; @Mock BigQueryResourceProvider bigQueryResourceProviderMock; - SimpleTable exampleTableInDefaultProject; - - SimpleTable exampleTableInNonDefaultProject; - - SimpleTable replacementTableInDefaultProject; - - FunctionInfo exampleFunction = - FunctionInfo.newBuilder() - .setNamePath(ImmutableList.of(testProjectId, "dataset", "examplefunction")) - .setGroup("UDF") - .setMode(Mode.SCALAR) - .setSignatures( - ImmutableList.of( - new FunctionSignature( - new FunctionArgumentType(TypeFactory.createSimpleType(TypeKind.TYPE_STRING)), - ImmutableList.of(), - -1))) - .build(); TVFInfo exampleTVF = TVFInfo.newBuilder() @@ -91,106 +80,64 @@ public class BigQueryCatalogTest { void init() { this.bigQueryCatalog = new BigQueryCatalog(this.testProjectId, this.bigQueryResourceProviderMock); - this.setupExampleTables(); - } - - private void setupExampleTables() { - String exampleTableName = "BigQueryTable"; - exampleTableInDefaultProject = - new SimpleTable( - exampleTableName, - ImmutableList.of( - new SimpleColumn( - exampleTableName, "col1", TypeFactory.createSimpleType(TypeKind.TYPE_STRING)))); - exampleTableInDefaultProject.setFullName( - String.format("%s.dataset.%s", this.testProjectId, exampleTableName)); - - exampleTableInNonDefaultProject = - new SimpleTable( - exampleTableInDefaultProject.getName(), exampleTableInDefaultProject.getColumnList()); - exampleTableInNonDefaultProject.setFullName( - String.format("%s.dataset.%s", "another-project-id", exampleTableName)); - - replacementTableInDefaultProject = - new SimpleTable( - exampleTableName, - ImmutableList.of( - new SimpleColumn( - exampleTableName, "col1", TypeFactory.createSimpleType(TypeKind.TYPE_STRING)), - new SimpleColumn( - exampleTableName, "col2", TypeFactory.createSimpleType(TypeKind.TYPE_STRING)))); - replacementTableInDefaultProject.setFullName( - String.format("%s.dataset.%s", this.testProjectId, exampleTableName)); } - private List> buildPathsWhereResourceShouldBe(String resourceReference) { - BigQueryReference ref = BigQueryReference.from(this.testProjectId, resourceReference); - - List> fullyQualifiedPaths = + private SimpleTable buildExampleTable(String tableRef) { + SimpleTable table = new SimpleTable( + tableRef, ImmutableList.of( - ImmutableList.of(ref.getProjectId() + "." + ref.getDatasetId() + "." + ref.getResourceName()), - ImmutableList.of(ref.getProjectId(), ref.getDatasetId() + "." + ref.getResourceName()), - ImmutableList.of(ref.getProjectId() + "." + ref.getDatasetId(), ref.getResourceName()), - ImmutableList.of(ref.getProjectId(), ref.getDatasetId(), ref.getResourceName())); - - List> result = new ArrayList<>(fullyQualifiedPaths); - - if (ref.getProjectId().equals(this.testProjectId)) { - List> implicitProjectPaths = - ImmutableList.of( - ImmutableList.of(ref.getDatasetId() + "." + ref.getResourceName()), - ImmutableList.of(ref.getDatasetId(), ref.getResourceName())); - result.addAll(implicitProjectPaths); - } + new SimpleColumn( + tableRef, + "col1", + TypeFactory.createSimpleType(TypeKind.TYPE_STRING)))); + table.setFullName(tableRef); + return table; + } - return result; + private FunctionInfo buildExampleFunction(String functionRef) { + return FunctionInfo.newBuilder() + .setNamePath(ImmutableList.of(functionRef)) + .setGroup("UDF") + .setMode(Mode.SCALAR) + .setSignatures( + ImmutableList.of( + new FunctionSignature( + new FunctionArgumentType(TypeFactory.createSimpleType(TypeKind.TYPE_STRING)), + ImmutableList.of(), + -1))) + .build(); } - private Table assertTableExistsAtPaths(BigQueryCatalog catalog, List> tablePaths) { - Preconditions.checkNotNull(tablePaths, "Table paths cannot be null"); - Preconditions.checkElementIndex(0, tablePaths.size(), "Table paths cannot be empty"); + private Table assertTableExists(BigQueryCatalog catalog, String tableRef) { + Preconditions.checkNotNull(catalog, "Catalog cannot be null"); + Preconditions.checkNotNull(tableRef, "Table reference cannot be null"); SimpleCatalog underlyingCatalog = catalog.getZetaSQLCatalog(); + List tablePath = ImmutableList.of(tableRef); - Stream assertions = - tablePaths.stream() - .map( - tablePath -> - (() -> - assertDoesNotThrow( - () -> underlyingCatalog.findTable(tablePath), - String.format( - "Expected table to exist at path %s", - String.join(".", tablePath))))); - - assertAll(assertions); + assertDoesNotThrow( + () -> underlyingCatalog.findTable(tablePath), + String.format( + "Expected table %s to exist", tableRef)); try { - return underlyingCatalog.findTable(tablePaths.get(0)); + return underlyingCatalog.findTable(tablePath); } catch (NotFoundException e) { throw new AssertionError(e); } } - private void assertTableDoesNotExist(BigQueryCatalog catalog, List> tablePaths) { - Preconditions.checkNotNull(tablePaths, "Table paths cannot be null"); - Preconditions.checkElementIndex(0, tablePaths.size(), "Table paths cannot be empty"); + private void assertTableDoesNotExist(BigQueryCatalog catalog, String tableRef) { + Preconditions.checkNotNull(catalog, "Catalog cannot be null"); + Preconditions.checkNotNull(tableRef, "Table reference cannot be null"); SimpleCatalog underlyingCatalog = catalog.getZetaSQLCatalog(); + List tablePath = ImmutableList.of(tableRef); - Stream assertions = - tablePaths.stream() - .map( - tablePath -> - (() -> - assertThrows( - NotFoundException.class, - () -> underlyingCatalog.findTable(tablePath), - String.format( - "Expected table to not exist at path %s", - String.join(".", tablePath))))); - - assertAll(assertions); + assertThrows( + NotFoundException.class, + () -> underlyingCatalog.findTable(tablePath), + String.format("Expected table %s to not exist", tableRef)); } @Test @@ -219,15 +166,8 @@ void testCatalogSupportsBigQueryTypeNames() { @Test void testRegisterInvalidTableName() { - String tableName = "TableName"; String invalidFullName = "An.Invalid.BQ.Reference"; - SimpleTable table = - new SimpleTable( - tableName, - ImmutableList.of( - new SimpleColumn( - tableName, "column", TypeFactory.createSimpleType(TypeKind.TYPE_STRING)))); - table.setFullName(invalidFullName); + SimpleTable table = buildExampleTable(invalidFullName); when(this.bigQueryResourceProviderMock.getTables(anyString(), anyList())) .thenReturn(ImmutableList.of(table)); @@ -247,88 +187,102 @@ void testRegisterInvalidTableName() { @Test void testRegisterTableFromDefaultProject() { - this.bigQueryCatalog.register( - exampleTableInDefaultProject, CreateMode.CREATE_DEFAULT, CreateScope.CREATE_DEFAULT_SCOPE); + SimpleTable table = + new SimpleTable( + "example", + ImmutableList.of( + new SimpleColumn( + "example", + "col1", + TypeFactory.createSimpleType(TypeKind.TYPE_STRING)))); + table.setFullName(String.format("%s.dataset.example", this.testProjectId)); - List> pathsWhereTableShouldBe = - this.buildPathsWhereResourceShouldBe(exampleTableInDefaultProject.getFullName()); + this.bigQueryCatalog.register( + table, CreateMode.CREATE_DEFAULT, CreateScope.CREATE_DEFAULT_SCOPE); - Table foundTable = assertTableExistsAtPaths(this.bigQueryCatalog, pathsWhereTableShouldBe); + Table foundTable = assertTableExists(bigQueryCatalog, table.getFullName()); + assertTableExists(bigQueryCatalog, "dataset.example"); assertTrue( - CatalogTestUtils.tableColumnsEqual(exampleTableInDefaultProject, foundTable), + CatalogTestUtils.tableColumnsEqual(table, foundTable), "Expected table created in Catalog to be equal to the original"); } @Test void testRegisterTableFromNonDefaultProject() { - this.bigQueryCatalog.register( - exampleTableInNonDefaultProject, - CreateMode.CREATE_DEFAULT, - CreateScope.CREATE_DEFAULT_SCOPE); - - List> pathsWhereTableShouldBe = - this.buildPathsWhereResourceShouldBe(exampleTableInNonDefaultProject.getFullName()); - - assertTableExistsAtPaths(this.bigQueryCatalog, pathsWhereTableShouldBe); + SimpleTable table = + new SimpleTable( + "example", + ImmutableList.of( + new SimpleColumn( + "example", + "col1", + TypeFactory.createSimpleType(TypeKind.TYPE_STRING)))); + table.setFullName("another-project-id.dataset.example"); - BigQueryReference ref = - BigQueryReference.from(this.testProjectId, exampleTableInNonDefaultProject.getFullName()); - List pathWhereTableShouldNotBe = ImmutableList.of(ref.getDatasetId(), ref.getResourceName()); + this.bigQueryCatalog.register( + table, CreateMode.CREATE_DEFAULT, CreateScope.CREATE_DEFAULT_SCOPE); - assertThrows( - NotFoundException.class, - () -> this.bigQueryCatalog.getZetaSQLCatalog().findTable(pathWhereTableShouldNotBe), - "Expected table not in default project to not be available at DATASET.TABLE path"); + assertTableExists(bigQueryCatalog, table.getFullName()); + assertTableDoesNotExist(bigQueryCatalog, "dataset.example"); } @Test void testRemoveTable() { - this.bigQueryCatalog.register( - exampleTableInDefaultProject, CreateMode.CREATE_DEFAULT, CreateScope.CREATE_DEFAULT_SCOPE); + String tableRef = String.format("%s.dataset.example", testProjectId); + SimpleTable table = buildExampleTable(tableRef); - List> pathsWhereTableShouldBe = - this.buildPathsWhereResourceShouldBe(exampleTableInDefaultProject.getFullName()); + bigQueryCatalog.register( + table, CreateMode.CREATE_DEFAULT, CreateScope.CREATE_DEFAULT_SCOPE); - assertTableExistsAtPaths(this.bigQueryCatalog, pathsWhereTableShouldBe); + assertTableExists(bigQueryCatalog, tableRef); + assertTableExists(bigQueryCatalog, "dataset.example"); - this.bigQueryCatalog.removeTable(exampleTableInDefaultProject.getFullName()); + this.bigQueryCatalog.removeTable(tableRef); - assertTableDoesNotExist(this.bigQueryCatalog, pathsWhereTableShouldBe); + assertTableDoesNotExist(bigQueryCatalog, tableRef); + assertTableDoesNotExist(bigQueryCatalog, "dataset.example"); } @Test void testReplaceTable() { - this.bigQueryCatalog.register( - exampleTableInDefaultProject, CreateMode.CREATE_DEFAULT, CreateScope.CREATE_DEFAULT_SCOPE); + String tableRef = String.format("%s.dataset.example", testProjectId); + SimpleTable table = buildExampleTable(tableRef); + + SimpleTable replacement = new SimpleTable(tableRef, ImmutableList.of( + new SimpleColumn(tableRef, "newCol", + TypeFactory.createSimpleType(TypeKind.TYPE_STRING)))); + + bigQueryCatalog.register( + table, CreateMode.CREATE_DEFAULT, CreateScope.CREATE_DEFAULT_SCOPE); // Replace the table and validate it has been replaced - this.bigQueryCatalog.register( - replacementTableInDefaultProject, + bigQueryCatalog.register( + replacement, CreateMode.CREATE_OR_REPLACE, CreateScope.CREATE_DEFAULT_SCOPE); - List> pathsWhereTableShouldBe = - this.buildPathsWhereResourceShouldBe(replacementTableInDefaultProject.getFullName()); - - Table foundTable = assertTableExistsAtPaths(this.bigQueryCatalog, pathsWhereTableShouldBe); + Table foundTable = assertTableExists(this.bigQueryCatalog, tableRef); assertTrue( - CatalogTestUtils.tableColumnsEqual(replacementTableInDefaultProject, foundTable), - "Expected table created in Catalog to be equal to the original"); + CatalogTestUtils.tableColumnsEqual(replacement, foundTable), + "Expected table to have been replaced"); } @Test void testTableAlreadyExists() { + String tableRef = String.format("%s.dataset.example", testProjectId); + SimpleTable table = buildExampleTable(tableRef); + this.bigQueryCatalog.register( - exampleTableInDefaultProject, CreateMode.CREATE_DEFAULT, CreateScope.CREATE_DEFAULT_SCOPE); + table, CreateMode.CREATE_DEFAULT, CreateScope.CREATE_DEFAULT_SCOPE); // Try to replace the table without using CreateMode.CREATE_OR_REPLACE - Assertions.assertThrows( + assertThrows( CatalogResourceAlreadyExists.class, () -> this.bigQueryCatalog.register( - replacementTableInDefaultProject, + table, CreateMode.CREATE_DEFAULT, CreateScope.CREATE_DEFAULT_SCOPE), "Expected fail creating table that already exists"); @@ -336,68 +290,78 @@ void testTableAlreadyExists() { @Test void testAddTablesByName() { + String tableRef = String.format("%s.dataset.example", testProjectId); + SimpleTable table = buildExampleTable(tableRef); + // When BigQueryResourceProvider.getTables() is called, return the test table when(bigQueryResourceProviderMock.getTables(anyString(), anyList())) - .thenReturn(ImmutableList.of(exampleTableInDefaultProject)); + .thenReturn(ImmutableList.of(table)); // Add the tables by name - bigQueryCatalog.addTables(ImmutableList.of(exampleTableInDefaultProject.getFullName())); + bigQueryCatalog.addTables(ImmutableList.of(table.getFullName())); // Verify the BigQueryCatalog got the tables from the BigQueryResourceProvider - verify(bigQueryResourceProviderMock, times(1)).getTables(anyString(), anyList()); + verify(bigQueryResourceProviderMock, times(1)) + .getTables(anyString(), anyList()); // Verify the test table was added to the catalog - List> pathsWhereTableShouldBe = - this.buildPathsWhereResourceShouldBe(exampleTableInDefaultProject.getFullName()); - assertTableExistsAtPaths(bigQueryCatalog, pathsWhereTableShouldBe); + assertTableExists(bigQueryCatalog, tableRef); } @Test void testAddAllTablesInDataset() { + String tableRef = String.format("%s.dataset.example", testProjectId); + SimpleTable table = buildExampleTable(tableRef); + when(bigQueryResourceProviderMock.getAllTablesInDataset(anyString(), anyString())) - .thenReturn(ImmutableList.of(exampleTableInDefaultProject)); + .thenReturn(ImmutableList.of(table)); bigQueryCatalog.addAllTablesInDataset(testProjectId, "dataset"); - verify(bigQueryResourceProviderMock, times(1)).getAllTablesInDataset(anyString(), anyString()); + verify(bigQueryResourceProviderMock, times(1)) + .getAllTablesInDataset(anyString(), anyString()); - List> pathsWhereTableShouldBe = - this.buildPathsWhereResourceShouldBe(exampleTableInDefaultProject.getFullName()); - assertTableExistsAtPaths(bigQueryCatalog, pathsWhereTableShouldBe); + assertTableExists(bigQueryCatalog, tableRef); } @Test void testAddAllTablesInProject() { + String tableRef = String.format("%s.dataset.example", testProjectId); + SimpleTable table = buildExampleTable(tableRef); + when(bigQueryResourceProviderMock.getAllTablesInProject(anyString())) - .thenReturn(ImmutableList.of(exampleTableInDefaultProject)); + .thenReturn(ImmutableList.of(table)); bigQueryCatalog.addAllTablesInProject(testProjectId); - verify(bigQueryResourceProviderMock, times(1)).getAllTablesInProject(anyString()); + verify(bigQueryResourceProviderMock, times(1)) + .getAllTablesInProject(anyString()); - List> pathsWhereTableShouldBe = - this.buildPathsWhereResourceShouldBe(exampleTableInDefaultProject.getFullName()); - Table foundTable = assertTableExistsAtPaths(bigQueryCatalog, pathsWhereTableShouldBe); + assertTableExists(bigQueryCatalog, tableRef); } @Test void testRegisterFunction() { + String functionRef = String.format("%s.dataset.function", this.testProjectId); + FunctionInfo function = buildExampleFunction(functionRef); + this.bigQueryCatalog.register( - exampleFunction, CreateMode.CREATE_DEFAULT, CreateScope.CREATE_DEFAULT_SCOPE); + function, CreateMode.CREATE_DEFAULT, CreateScope.CREATE_DEFAULT_SCOPE); SimpleCatalog underlyingCatalog = this.bigQueryCatalog.getZetaSQLCatalog(); assertDoesNotThrow( - () -> underlyingCatalog.findFunction(exampleFunction.getNamePath()), + () -> underlyingCatalog.findFunction(function.getNamePath()), String.format( "Expected function to exist at path %s", - String.join(".", exampleFunction.getNamePath()))); + String.join(".", function.getNamePath()))); } @Test void testInferFunctionReturnType() { + String functionRef = String.format("%s.dataset.function", this.testProjectId); FunctionInfo functionWithUnknownReturnType = FunctionInfo.newBuilder() - .setNamePath(ImmutableList.of(testProjectId, "dataset", "function")) + .setNamePath(ImmutableList.of(functionRef)) .setGroup("UDF") .setMode(Mode.SCALAR) .setSignatures( diff --git a/zetasql-toolkit-core/pom.xml b/zetasql-toolkit-core/pom.xml index 73c153e..63e5ca1 100644 --- a/zetasql-toolkit-core/pom.xml +++ b/zetasql-toolkit-core/pom.xml @@ -22,12 +22,12 @@ com.google.zetasql.toolkit zetasql-toolkit - 0.4.1 + 0.5.0 ../pom.xml zetasql-toolkit-core - 0.4.1 + 0.5.0 ${project.groupId}:${project.artifactId} @@ -105,14 +105,6 @@ com.google.cloud google-cloud-spanner - - org.slf4j - slf4j-api - - - org.slf4j - slf4j-simple - org.junit.jupiter junit-jupiter-engine diff --git a/zetasql-toolkit-core/src/main/antlr4/com/google/zetasql/toolkit/catalog/typeparser/ZetaSQLTypeGrammar.g4 b/zetasql-toolkit-core/src/main/antlr4/com/google/zetasql/toolkit/catalog/typeparser/ZetaSQLTypeGrammar.g4 index 918766a..333055d 100644 --- a/zetasql-toolkit-core/src/main/antlr4/com/google/zetasql/toolkit/catalog/typeparser/ZetaSQLTypeGrammar.g4 +++ b/zetasql-toolkit-core/src/main/antlr4/com/google/zetasql/toolkit/catalog/typeparser/ZetaSQLTypeGrammar.g4 @@ -34,24 +34,53 @@ BASIC_TYPE: STRING | GEOGRAPHY | JSON; -STRING: 'STRING'; -BYTES: 'BYTES'; -INT32: 'INT32'; -INT64: 'INT64'; -UINT32: 'UINT32'; -UINT64: 'UINT64'; -FLOAT64: 'FLOAT64'; -DECIMAL: 'DECIMAL'; -NUMERIC: 'NUMERIC'; -BIGNUMERIC: 'BIGNUMERIC'; -INTERVAL: 'INTERVAL'; -BOOL: 'BOOL'; -TIMESTAMP: 'TIMESTAMP'; -DATE: 'DATE'; -TIME: 'TIME'; -DATETIME: 'DATETIME'; -GEOGRAPHY: 'GEOGRAPHY'; -JSON: 'JSON'; +STRING: S T R I N G; +BYTES: B Y T E S; +INT32: I N T THIRTYTWO; +INT64: I N T SIXTYFOUR; +UINT32: U I N T THIRTYTWO; +UINT64: U I N T SIXTYFOUR; +FLOAT64: F L O A T SIXTYFOUR; +DECIMAL: D E C I M A L; +NUMERIC: N U M E R I C; +BIGNUMERIC: B I G N U M E R I C; +INTERVAL: I N T E R V A L; +BOOL: B O O L; +TIMESTAMP: T I M E S T A M P; +DATE: D A T E; +TIME: T I M E; +DATETIME: D A T E T I M E; +GEOGRAPHY: G E O G R A P H Y; +JSON: J S O N; IDENTIFIER: [A-Za-z_][A-Za-z0-9_]*; NUMBER: [1-9][0-9]*; + +fragment THIRTYTWO: '32'; +fragment SIXTYFOUR: '64'; +fragment A: [aA]; +fragment B: [bB]; +fragment C: [cC]; +fragment D: [dD]; +fragment E: [eE]; +fragment F: [fF]; +fragment G: [gG]; +fragment H: [hH]; +fragment I: [iI]; +fragment J: [jJ]; +fragment K: [kK]; +fragment L: [lL]; +fragment M: [mM]; +fragment N: [nN]; +fragment O: [oO]; +fragment P: [pP]; +fragment Q: [qQ]; +fragment R: [rR]; +fragment S: [sS]; +fragment T: [tT]; +fragment U: [uU]; +fragment V: [vV]; +fragment W: [wW]; +fragment X: [xX]; +fragment Y: [yY]; +fragment Z: [zZ]; diff --git a/zetasql-toolkit-core/src/main/java/com/google/zetasql/SimpleCatalogUtil.java b/zetasql-toolkit-core/src/main/java/com/google/zetasql/SimpleCatalogUtil.java new file mode 100644 index 0000000..15eefe1 --- /dev/null +++ b/zetasql-toolkit-core/src/main/java/com/google/zetasql/SimpleCatalogUtil.java @@ -0,0 +1,32 @@ +/* + * Copyright 2023 Google LLC All Rights Reserved + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.zetasql; + +/** + * Only public to make it possible to copy the SimpleCatalog + */ +public class SimpleCatalogUtil { + public static SimpleCatalog copyCatalog(SimpleCatalog sourceCatalog) { + // Simply serializes and deserializes the source catalog to create a copy. + // This is the most reliable way of creating a copy of a SimpleCatalog, + // as the SimpleCatalog's public interface does not expose enough of the internal + // structures to create an accurate copy. + FileDescriptorSetsBuilder fileDescriptorSetsBuilder = new FileDescriptorSetsBuilder(); + SimpleCatalogProtos.SimpleCatalogProto serialized = sourceCatalog.serialize(fileDescriptorSetsBuilder); + return SimpleCatalog.deserialize(serialized, fileDescriptorSetsBuilder.getDescriptorPools()); + } +} diff --git a/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/StatementRewriter.java b/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/StatementRewriter.java new file mode 100644 index 0000000..2d66780 --- /dev/null +++ b/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/StatementRewriter.java @@ -0,0 +1,538 @@ +/* + * Copyright 2023 Google LLC All Rights Reserved + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.zetasql.toolkit; + +import com.google.zetasql.parser.ASTNode; +import com.google.zetasql.parser.ASTNodes.ASTAlterAllRowAccessPoliciesStatement; +import com.google.zetasql.parser.ASTNodes.ASTAlterApproxViewStatement; +import com.google.zetasql.parser.ASTNodes.ASTAlterDatabaseStatement; +import com.google.zetasql.parser.ASTNodes.ASTAlterEntityStatement; +import com.google.zetasql.parser.ASTNodes.ASTAlterMaterializedViewStatement; +import com.google.zetasql.parser.ASTNodes.ASTAlterModelStatement; +import com.google.zetasql.parser.ASTNodes.ASTAlterPrivilegeRestrictionStatement; +import com.google.zetasql.parser.ASTNodes.ASTAlterRowAccessPolicyStatement; +import com.google.zetasql.parser.ASTNodes.ASTAlterSchemaStatement; +import com.google.zetasql.parser.ASTNodes.ASTAlterTableStatement; +import com.google.zetasql.parser.ASTNodes.ASTAlterViewStatement; +import com.google.zetasql.parser.ASTNodes.ASTAuxLoadDataStatement; +import com.google.zetasql.parser.ASTNodes.ASTCallStatement; +import com.google.zetasql.parser.ASTNodes.ASTCloneDataSource; +import com.google.zetasql.parser.ASTNodes.ASTCloneDataStatement; +import com.google.zetasql.parser.ASTNodes.ASTConnectionClause; +import com.google.zetasql.parser.ASTNodes.ASTCopyDataSource; +import com.google.zetasql.parser.ASTNodes.ASTCreateApproxViewStatement; +import com.google.zetasql.parser.ASTNodes.ASTCreateConstantStatement; +import com.google.zetasql.parser.ASTNodes.ASTCreateDatabaseStatement; +import com.google.zetasql.parser.ASTNodes.ASTCreateEntityStatement; +import com.google.zetasql.parser.ASTNodes.ASTCreateExternalTableStatement; +import com.google.zetasql.parser.ASTNodes.ASTCreateIndexStatement; +import com.google.zetasql.parser.ASTNodes.ASTCreateMaterializedViewStatement; +import com.google.zetasql.parser.ASTNodes.ASTCreateModelStatement; +import com.google.zetasql.parser.ASTNodes.ASTCreatePrivilegeRestrictionStatement; +import com.google.zetasql.parser.ASTNodes.ASTCreateProcedureStatement; +import com.google.zetasql.parser.ASTNodes.ASTCreateRowAccessPolicyStatement; +import com.google.zetasql.parser.ASTNodes.ASTCreateSchemaStatement; +import com.google.zetasql.parser.ASTNodes.ASTCreateSnapshotStatement; +import com.google.zetasql.parser.ASTNodes.ASTCreateSnapshotTableStatement; +import com.google.zetasql.parser.ASTNodes.ASTCreateTableStatement; +import com.google.zetasql.parser.ASTNodes.ASTCreateViewStatement; +import com.google.zetasql.parser.ASTNodes.ASTDefineTableStatement; +import com.google.zetasql.parser.ASTNodes.ASTDescribeStatement; +import com.google.zetasql.parser.ASTNodes.ASTDropAllRowAccessPoliciesStatement; +import com.google.zetasql.parser.ASTNodes.ASTDropEntityStatement; +import com.google.zetasql.parser.ASTNodes.ASTDropFunctionStatement; +import com.google.zetasql.parser.ASTNodes.ASTDropMaterializedViewStatement; +import com.google.zetasql.parser.ASTNodes.ASTDropPrivilegeRestrictionStatement; +import com.google.zetasql.parser.ASTNodes.ASTDropRowAccessPolicyStatement; +import com.google.zetasql.parser.ASTNodes.ASTDropSearchIndexStatement; +import com.google.zetasql.parser.ASTNodes.ASTDropSnapshotTableStatement; +import com.google.zetasql.parser.ASTNodes.ASTDropStatement; +import com.google.zetasql.parser.ASTNodes.ASTDropTableFunctionStatement; +import com.google.zetasql.parser.ASTNodes.ASTDropVectorIndexStatement; +import com.google.zetasql.parser.ASTNodes.ASTExportMetadataStatement; +import com.google.zetasql.parser.ASTNodes.ASTExportModelStatement; +import com.google.zetasql.parser.ASTNodes.ASTForeignKeyReference; +import com.google.zetasql.parser.ASTNodes.ASTFunctionCall; +import com.google.zetasql.parser.ASTNodes.ASTFunctionDeclaration; +import com.google.zetasql.parser.ASTNodes.ASTGrantStatement; +import com.google.zetasql.parser.ASTNodes.ASTIdentifier; +import com.google.zetasql.parser.ASTNodes.ASTImportStatement; +import com.google.zetasql.parser.ASTNodes.ASTMergeStatement; +import com.google.zetasql.parser.ASTNodes.ASTModelClause; +import com.google.zetasql.parser.ASTNodes.ASTModuleStatement; +import com.google.zetasql.parser.ASTNodes.ASTPathExpression; +import com.google.zetasql.parser.ASTNodes.ASTRenameStatement; +import com.google.zetasql.parser.ASTNodes.ASTRenameToClause; +import com.google.zetasql.parser.ASTNodes.ASTRevokeStatement; +import com.google.zetasql.parser.ASTNodes.ASTSequenceArg; +import com.google.zetasql.parser.ASTNodes.ASTShowStatement; +import com.google.zetasql.parser.ASTNodes.ASTSpannerInterleaveClause; +import com.google.zetasql.parser.ASTNodes.ASTStatement; +import com.google.zetasql.parser.ASTNodes.ASTSystemVariableExpr; +import com.google.zetasql.parser.ASTNodes.ASTTVF; +import com.google.zetasql.parser.ASTNodes.ASTTableAndColumnInfo; +import com.google.zetasql.parser.ASTNodes.ASTTableClause; +import com.google.zetasql.parser.ASTNodes.ASTTablePathExpression; +import com.google.zetasql.parser.ASTNodes.ASTTruncateStatement; +import com.google.zetasql.parser.ASTNodes.ASTUndropStatement; +import com.google.zetasql.parser.ParseTreeVisitor; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.List; +import java.util.stream.Collectors; + +/** + * Implements query-level rewrites based on the parsed tree. It allows modifying the query text + * after parsing but before analyzing. + * + *

These rewrites are done at the query level because the parsed tree is immutable and can't + * be modified itself. + */ +class StatementRewriter { + + /** + * Rewrites the query text to ensure all resource name paths present in a parsed statement from + * said query are quoted entirely. For example, it rewrites "FROM project.dataset.table" to "FROM + * `project.dataset.table`". + * + *

To do so; it finds all {@link ASTPathExpression} nodes in the parse tree refer to a + * resource (i.e. tables, functions, etc.), builds their fully quoted representation and replaces + * their original text in the query with their quoted representation. + * + * @param query The original query text the statement was parsed from + * @param parsedStatement The parsed statement for which to rewrite name paths + * @return The rewritten version of the query + */ + public static String quoteNamePaths(String query, ASTStatement parsedStatement) { + List pathExpressions = + getResourcePathExpressionFromParseTree(parsedStatement).stream() + .sorted(Comparator.comparing(expression -> expression.getParseLocationRange().start())) + .collect(Collectors.toList()); + + StringBuilder builder = new StringBuilder(query); + + int rewritingOffset = 0; + for (ASTPathExpression pathExpression : pathExpressions) { + int startPosition = pathExpression.getParseLocationRange().start(); + int endPosition = pathExpression.getParseLocationRange().end(); + + String originalNamePath = query.substring(startPosition, endPosition); + String quotedNamePath = buildQuotedNamePath(pathExpression); + + int rewriteStartPosition = startPosition + rewritingOffset; + int rewriteEndPosition = endPosition + rewritingOffset; + + builder.replace(rewriteStartPosition, rewriteEndPosition, quotedNamePath); + + rewritingOffset += (quotedNamePath.length() - originalNamePath.length()); + } + + return builder.toString(); + } + + /** + * Returns the fully quoted string representation of an {@link ASTPathExpression}. + * + * @param pathExpression The path expression for which to build the fully quoted + * representation + * @return The fully quoted representation of the path expression + */ + private static String buildQuotedNamePath(ASTPathExpression pathExpression) { + String fullName = pathExpression.getNames() + .stream() + .map(ASTIdentifier::getIdString) + .collect(Collectors.joining(".")); + return "`" + fullName + "`"; + } + + /** + * Extracts all {@link ASTPathExpression} nodes that refer to a resource + * (i.e. tables, functions, etc.) from a parse tree + * + * @param tree The parse tree + * @return The list of all {@link ASTPathExpression} in the tree that refer to a resource + */ + private static List getResourcePathExpressionFromParseTree(ASTNode tree) { + ArrayList result = new ArrayList<>(); + + tree.accept(new ParseTreeVisitor() { + + public void visit(ASTTablePathExpression node) { + result.add(node.getPathExpr()); + } + + public void visit(ASTFunctionCall node) { + result.add(node.getFunction()); + super.visit(node); + } + + public void visit(ASTSequenceArg node) { + result.add(node.getSequencePath()); + } + + public void visit(ASTDescribeStatement node) { + result.add(node.getName()); + if (node.getOptionalFromName() != null) { + result.add(node.getOptionalFromName()); + } + super.visit(node); + } + + public void visit(ASTShowStatement node) { + result.add(node.getOptionalName()); + super.visit(node); + } + + public void visit(ASTDropEntityStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTDropFunctionStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTDropTableFunctionStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTDropAllRowAccessPoliciesStatement node) { + result.add(node.getTableName()); + super.visit(node); + } + + public void visit(ASTDropMaterializedViewStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTDropSnapshotTableStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTDropSearchIndexStatement node) { + result.add(node.getName()); + result.add(node.getTableName()); + super.visit(node); + } + + public void visit(ASTDropVectorIndexStatement node) { + result.add(node.getName()); + result.add(node.getTableName()); + super.visit(node); + } + + public void visit(ASTRenameStatement node) { + result.add(node.getOldName()); + result.add(node.getNewName()); + super.visit(node); + } + + public void visit(ASTImportStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTModuleStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTSystemVariableExpr node) { + result.add(node.getPath()); + } + + public void visit(ASTFunctionDeclaration node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTTVF node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTTableClause node) { + result.add(node.getTablePath()); + super.visit(node); + } + + public void visit(ASTModelClause node) { + result.add(node.getModelPath()); + } + + public void visit(ASTConnectionClause node) { + result.add(node.getConnectionPath()); + } + + public void visit(ASTCloneDataSource node) { + result.add(node.getPathExpr()); + super.visit(node); + } + + public void visit(ASTCopyDataSource node) { + result.add(node.getPathExpr()); + super.visit(node); + } + + public void visit(ASTCloneDataStatement node) { + result.add(node.getTargetPath()); + super.visit(node); + } + + public void visit(ASTCreateConstantStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTCreateDatabaseStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTCreateProcedureStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTCreateSchemaStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTCreateModelStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTCreateIndexStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTExportModelStatement node) { + result.add(node.getModelNamePath()); + super.visit(node); + } + + public void visit(ASTExportMetadataStatement node) { + result.add(node.getNamePath()); + super.visit(node); + } + + public void visit(ASTCallStatement node) { + result.add(node.getProcedureName()); + super.visit(node); + } + + public void visit(ASTDefineTableStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTCreateSnapshotStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTCreateSnapshotTableStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTTableAndColumnInfo node) { + result.add(node.getTableName()); + super.visit(node); + } + + public void visit(ASTTruncateStatement node) { + result.add(node.getTargetPath()); + super.visit(node); + } + + public void visit(ASTMergeStatement node) { + result.add(node.getTargetPath()); + super.visit(node); + } + + public void visit(ASTGrantStatement node) { + result.add(node.getTargetPath()); + super.visit(node); + } + + public void visit(ASTRevokeStatement node) { + result.add(node.getTargetPath()); + super.visit(node); + } + + public void visit(ASTRenameToClause node) { + result.add(node.getNewName()); + super.visit(node); + } + + public void visit(ASTAlterAllRowAccessPoliciesStatement node) { + result.add(node.getTableNamePath()); + super.visit(node); + } + + public void visit(ASTForeignKeyReference node) { + result.add(node.getTableName()); + super.visit(node); + } + + public void visit(ASTCreateEntityStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTDropPrivilegeRestrictionStatement node) { + result.add(node.getNamePath()); + super.visit(node); + } + + public void visit(ASTDropRowAccessPolicyStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTCreatePrivilegeRestrictionStatement node) { + result.add(node.getNamePath()); + super.visit(node); + } + + public void visit(ASTCreateRowAccessPolicyStatement node) { + result.add(node.getName()); + result.add(node.getTargetPath()); + super.visit(node); + } + + public void visit(ASTDropStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTCreateTableStatement node) { + result.add(node.getName()); + if (node.getLikeTableName() != null) { + result.add(node.getLikeTableName()); + } + super.visit(node); + } + + public void visit(ASTCreateExternalTableStatement node) { + result.add(node.getName()); + if (node.getLikeTableName() != null) { + result.add(node.getLikeTableName()); + } + super.visit(node); + } + + public void visit(ASTAuxLoadDataStatement node) { + result.add(node.getName()); + if (node.getLikeTableName() != null) { + result.add(node.getLikeTableName()); + } + super.visit(node); + } + + public void visit(ASTCreateViewStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTCreateApproxViewStatement node) { + result.add(node.getName()); + super.visit(node); + } + + public void visit(ASTCreateMaterializedViewStatement node) { + result.add(node.getName()); + if (node.getReplicaSource() != null) { + result.add(node.getReplicaSource()); + } + super.visit(node); + } + + public void visit(ASTAlterDatabaseStatement node) { + result.add(node.getPath()); + super.visit(node); + } + + public void visit(ASTAlterSchemaStatement node) { + result.add(node.getPath()); + super.visit(node); + } + + public void visit(ASTAlterTableStatement node) { + result.add(node.getPath()); + super.visit(node); + } + + public void visit(ASTAlterViewStatement node) { + result.add(node.getPath()); + super.visit(node); + } + + public void visit(ASTAlterMaterializedViewStatement node) { + result.add(node.getPath()); + super.visit(node); + } + + public void visit(ASTAlterApproxViewStatement node) { + result.add(node.getPath()); + super.visit(node); + } + + public void visit(ASTAlterModelStatement node) { + result.add(node.getPath()); + super.visit(node); + } + + public void visit(ASTAlterPrivilegeRestrictionStatement node) { + result.add(node.getPath()); + super.visit(node); + } + + public void visit(ASTAlterRowAccessPolicyStatement node) { + result.add(node.getPath()); + super.visit(node); + } + + public void visit(ASTAlterEntityStatement node) { + result.add(node.getPath()); + super.visit(node); + } + + public void visit(ASTSpannerInterleaveClause node) { + result.add(node.getTableName()); + super.visit(node); + } + + public void visit(ASTUndropStatement node) { + result.add(node.getName()); + super.visit(node); + } + + }); + + return result; + } + +} diff --git a/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/ZetaSQLToolkitAnalyzer.java b/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/ZetaSQLToolkitAnalyzer.java index 077a6f7..b6b95aa 100644 --- a/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/ZetaSQLToolkitAnalyzer.java +++ b/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/ZetaSQLToolkitAnalyzer.java @@ -182,7 +182,7 @@ public boolean hasNext() { */ @Override public AnalyzedStatement next() { - int startLocation = parseResumeLocation.getBytePosition(); + int startPosition = parseResumeLocation.getBytePosition(); ASTStatement parsedStatement = parseNextStatement(parseResumeLocation); @@ -206,9 +206,12 @@ public AnalyzedStatement next() { return new AnalyzedStatement(parsedStatement, Optional.empty()); } - parseResumeLocation.setBytePosition(startLocation); + String rewrittenQuery = + StatementRewriter.quoteNamePaths(query, parsedStatement); + ParseResumeLocation analysisParseResumeLocation = new ParseResumeLocation(rewrittenQuery); + analysisParseResumeLocation.setBytePosition(startPosition); - ResolvedStatement resolvedStatement = analyzeNextStatement(parseResumeLocation); + ResolvedStatement resolvedStatement = analyzeNextStatement(analysisParseResumeLocation); this.applyCatalogMutation(resolvedStatement); diff --git a/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/CatalogOperations.java b/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/CatalogOperations.java index 3a60225..ba8afc3 100644 --- a/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/CatalogOperations.java +++ b/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/CatalogOperations.java @@ -23,6 +23,7 @@ import com.google.zetasql.TableValuedFunction.FixedOutputSchemaTVF; import com.google.zetasql.resolvedast.ResolvedCreateStatementEnums.CreateMode; import com.google.zetasql.toolkit.catalog.exceptions.CatalogResourceAlreadyExists; +import com.google.zetasql.toolkit.catalog.exceptions.CatalogResourceDoesNotExist; import java.util.Arrays; import java.util.List; import java.util.Optional; @@ -38,13 +39,6 @@ * */ public class CatalogOperations { - // TODO: Probably come up with an abstraction to reduce code repetition in this class. - // This implementation has a lot of repeated code; namely in methods like - // validate[Resource]DoesNotExist(), delete[Resource]FromCatalog() and - // create[Resource]InCatalog(). - // Because of the slightly different ways the SimpleCatalog handles naming for different types of - // resources, - // avoiding that repetition is not very straightforward. private CatalogOperations() {} @@ -101,21 +95,11 @@ private static boolean functionExists(SimpleCatalog catalog, String fullName) { .anyMatch(fullNameWithoutGroup::equalsIgnoreCase); } - /** Returns true if the Function exists in the SimpleCatalog */ - private static boolean functionExists(SimpleCatalog catalog, Function function) { - return functionExists(catalog, function.getFullName(false)); - } - /** Returns true if the TVF named tvfName exists in the SimpleCatalog */ private static boolean tvfExists(SimpleCatalog catalog, String tvfName) { return catalog.getTVFNameList().contains(tvfName.toLowerCase()); } - /** Returns true if the TVF exists in the SimpleCatalog */ - private static boolean tvfExists(SimpleCatalog catalog, TableValuedFunction tvf) { - return tvfExists(catalog, tvf.getName()); - } - /** Returns true if the named procedureName exists in the SimpleCatalog */ private static boolean procedureExists(SimpleCatalog catalog, String procedureName) { return catalog.getProcedureList().stream() @@ -123,11 +107,6 @@ private static boolean procedureExists(SimpleCatalog catalog, String procedureNa .anyMatch(name -> name.equalsIgnoreCase(procedureName)); } - /** Returns true if the Procedure exists in the SimpleCatalog */ - private static boolean procedureExists(SimpleCatalog catalog, Procedure procedure) { - return procedureExists(catalog, procedure.getName()); - } - /** * Gets the SimpleCatalog in which a resource should be created, based on the root catalog and the * resource path. @@ -154,336 +133,277 @@ private static SimpleCatalog getSubCatalogForResource( } /** - * Checks a table does not exist at any of the provided paths. + * Generic function for creating a resource in a {@link SimpleCatalog} * - * @param rootCatalog The catalog to look for tables in - * @param tablePaths The list of paths the table should not be in - * @param fullTableName The full name of the table we're looking for, used for error reporting - * @throws CatalogResourceAlreadyExists if a table exists at any of the provided paths + * @param nameInCatalog The name the resource will have in the catalog + * @param createMode The {@link CreateMode} to use + * @param resourceType The resource type name (e.g. "table", "function") + * @param alreadyExists Whether the resource already exists + * @param creator A {@link Runnable} that will create the resource in the catalog if run + * @param deleter A Runnable that will delete the resource from the catalog if run */ - private static void validateTableDoesNotExist( - SimpleCatalog rootCatalog, List> tablePaths, String fullTableName) { - for (List tablePath : tablePaths) { - if (tableExists(rootCatalog, tablePath)) { - String errorMessage = - String.format( - "Table %s already exists at path %s", fullTableName, tablePath.toString()); - throw new CatalogResourceAlreadyExists(fullTableName, errorMessage); - } + private static void createResource( + String nameInCatalog, CreateMode createMode, String resourceType, + boolean alreadyExists, Runnable creator, Runnable deleter + ) { + if (createMode.equals(CreateMode.CREATE_IF_NOT_EXISTS) && alreadyExists) { + return; } + + if (createMode.equals(CreateMode.CREATE_OR_REPLACE) && alreadyExists) { + deleter.run(); + } + + if (createMode.equals(CreateMode.CREATE_DEFAULT) && alreadyExists) { + String errorMessage = + String.format( + "%s %s already exists in catalog", resourceType, nameInCatalog); + throw new CatalogResourceAlreadyExists(nameInCatalog, errorMessage); + } + + creator.run(); } /** - * Deletes a table from the specified paths in a {@link SimpleCatalog} + * Deletes a table with the provided name from {@link SimpleCatalog} * - * @param rootCatalog The catalog from which to delete tables - * @param tablePaths The paths for the table that should be deleted + * @param catalog The catalog from which to delete tables + * @param name The name for the table in the catalog + * @throws CatalogResourceDoesNotExist if the table does not exist in the catalog */ - public static void deleteTableFromCatalog( - SimpleCatalog rootCatalog, List> tablePaths) { - for (List tablePath : tablePaths) { - String tableName = tablePath.get(tablePath.size() - 1); - SimpleCatalog catalog = getSubCatalogForResource(rootCatalog, tablePath); - - if (tableExists(catalog, tableName)) { - catalog.removeSimpleTable(tableName); - } + public static void deleteTableFromCatalog(SimpleCatalog catalog, String name) { + if (!tableExists(catalog, name)) { + String errorMessage = String.format("Tried to delete table which does not exist: %s", name); + throw new CatalogResourceDoesNotExist(name, errorMessage); } + + catalog.removeSimpleTable(name); } /** - * Creates a table in a SimpleCatalog using the provided paths and complying with the provided - * CreateMode. + * Creates a table in a {@link SimpleCatalog} using the provided paths and complying with + * the provided CreateMode. * - * @param rootCatalog The root SimpleCatalog in which to create the table. - * @param tablePaths The table paths to create the table at. If multiple paths are provided, - * multiple copies of the table will be registered in the catalog. - * @param fullTableName The full name of the table to create. - * @param columns The list of columns for the table + * @param catalog The catalog in which to create the table + * @param nameInCatalog The name under which the table will be registered in the catalog + * @param table The {@link SimpleTable} object representing the table * @param createMode The CreateMode to use * @throws CatalogResourceAlreadyExists if the table already exists at any of the provided paths * and CreateMode != CREATE_OR_REPLACE. */ public static void createTableInCatalog( - SimpleCatalog rootCatalog, - List> tablePaths, - String fullTableName, - List columns, + SimpleCatalog catalog, + String nameInCatalog, + SimpleTable table, CreateMode createMode) { - if (createMode.equals(CreateMode.CREATE_OR_REPLACE)) { - deleteTableFromCatalog(rootCatalog, tablePaths); - } - - if (createMode.equals(CreateMode.CREATE_DEFAULT)) { - validateTableDoesNotExist(rootCatalog, tablePaths, fullTableName); - } + boolean alreadyExists = tableExists(catalog, nameInCatalog); - SimpleTable table = new SimpleTable(fullTableName, columns); - table.setFullName(fullTableName); - - for (List tablePath : tablePaths) { - String tableName = tablePath.get(tablePath.size() - 1); - SimpleCatalog catalogForCreation = getSubCatalogForResource(rootCatalog, tablePath); - - if (!tableExists(catalogForCreation, tableName)) { - catalogForCreation.addSimpleTable(tableName, table); - } - } + createResource( + nameInCatalog, + createMode, + "Table", + alreadyExists, + /*creator=*/ () -> catalog.addSimpleTable(nameInCatalog, table), + /*deleter=*/ () -> deleteTableFromCatalog(catalog, nameInCatalog) + ); } /** - * Checks a function does not exist at any of the provided paths. + * Deletes a function with the provided name from the {@link SimpleCatalog} * - * @param rootCatalog The catalog to look for functions in - * @param functionPaths The list of paths the function should not be in - * @param functionFullName The full name of the function we're looking for, only used for error - * reporting - * @throws CatalogResourceAlreadyExists if a function exists at any of the provided paths + * @param catalog The catalog from which to delete the function + * @param fullName The full name of the function in the catalog + * @throws CatalogResourceDoesNotExist if the function does not exist in the catalog */ - private static void validateFunctionDoesNotExist( - SimpleCatalog rootCatalog, List> functionPaths, String functionFullName) { - for (List functionPath : functionPaths) { - String functionNameInCatalog = functionPath.get(functionPath.size() - 1); - - SimpleCatalog catalog = getSubCatalogForResource(rootCatalog, functionPath); + public static void deleteFunctionFromCatalog(SimpleCatalog catalog, String fullName) { + String fullNameWithoutGroup = removeGroupFromFunctionName(fullName); - if (functionExists(catalog, functionNameInCatalog)) { - throw new CatalogResourceAlreadyExists(functionFullName); - } - } - } + Optional fullNameToDelete = + catalog.getFunctionNameList().stream() + .filter( + name -> + removeGroupFromFunctionName(name).equalsIgnoreCase(fullNameWithoutGroup)) + .findFirst(); - /** - * Deletes a function from the specified paths in a {@link SimpleCatalog} - * - * @param rootCatalog The catalog from which to delete functions - * @param functionPaths The paths for the function that should be deleted - */ - public static void deleteFunctionFromCatalog( - SimpleCatalog rootCatalog, List> functionPaths) { - for (List functionPath : functionPaths) { - String functionNameInCatalog = functionPath.get(functionPath.size() - 1); - SimpleCatalog catalog = getSubCatalogForResource(rootCatalog, functionPath); - - if (functionExists(catalog, functionNameInCatalog)) { - Optional fullNameToDelete = - catalog.getFunctionNameList().stream() - .filter( - fullName -> - removeGroupFromFunctionName(fullName) - .equalsIgnoreCase(functionNameInCatalog)) - .findFirst(); - fullNameToDelete.ifPresent(catalog::removeFunction); - } + if (fullNameToDelete.isPresent()) { + catalog.removeFunction(fullNameToDelete.get()); + } else { + String errorMessage = String.format( + "Tried to delete function which does not exist: %s", fullName); + throw new CatalogResourceDoesNotExist(fullName, errorMessage); } } /** - * Creates a function in a SimpleCatalog using the provided paths and complying with the provided - * CreateMode. + * Creates a function in a {@link SimpleCatalog} using the provided paths and complying with the + * provided CreateMode. * - * @param rootCatalog The root SimpleCatalog in which to create the function. - * @param functionPaths The function paths to create the function at. If multiple paths are - * provided, multiple copies of the function will be registered in the catalog. - * @param functionInfo The FunctionInfo object representing the function that should be created + * @param catalog The catalog in which to create the function + * @param nameInCatalog The name under which the function will be registered in the catalog + * @param functionInfo The {@link FunctionInfo} object representing the function that + * should be created * @param createMode The CreateMode to use * @throws CatalogResourceAlreadyExists if the function already exists at any of the provided * paths and CreateMode != CREATE_OR_REPLACE. */ public static void createFunctionInCatalog( - SimpleCatalog rootCatalog, - List> functionPaths, + SimpleCatalog catalog, + String nameInCatalog, FunctionInfo functionInfo, CreateMode createMode) { - if (createMode.equals(CreateMode.CREATE_OR_REPLACE)) { - deleteFunctionFromCatalog(rootCatalog, functionPaths); - } - - if (createMode.equals(CreateMode.CREATE_DEFAULT)) { - validateFunctionDoesNotExist( - rootCatalog, functionPaths, String.join(".", functionInfo.getNamePath())); - } - - for (List functionPath : functionPaths) { - String functionName = functionPath.get(functionPath.size() - 1); - SimpleCatalog catalogForCreation = getSubCatalogForResource(rootCatalog, functionPath); - - Function finalFunction = - new Function( - ImmutableList.of(functionName), - functionInfo.getGroup(), - functionInfo.getMode(), - functionInfo.getSignatures()); - - if (!functionExists(catalogForCreation, finalFunction)) { - catalogForCreation.addFunction(finalFunction); - } - } + boolean alreadyExists = functionExists(catalog, nameInCatalog); + + Function function = + new Function( + ImmutableList.of(nameInCatalog), + functionInfo.getGroup(), + functionInfo.getMode(), + functionInfo.getSignatures()); + + createResource( + nameInCatalog, + createMode, + "Function", + alreadyExists, + /*creator=*/ () -> catalog.addFunction(function), + /*deleter=*/ () -> deleteFunctionFromCatalog(catalog, nameInCatalog) + ); } /** - * Checks a TVF does not exist at any of the provided paths. + * Deletes a TVF with the provided name from the {@link SimpleCatalog} * - * @param rootCatalog The catalog to look for functions in - * @param functionPaths The list of paths the function should not be in - * @param functionFullName The full name of the function we're looking for, only used for error - * reporting - * @throws CatalogResourceAlreadyExists if a function exists at any of the provided paths + * @param catalog The catalog from which to delete the function + * @param fullName The full name of the function in the catalog + * @throws CatalogResourceDoesNotExist if the function does not exist in the catalog */ - private static void validateTVFDoesNotExist( - SimpleCatalog rootCatalog, List> functionPaths, String functionFullName) { - for (List functionPath : functionPaths) { - String functionName = functionPath.get(functionPath.size() - 1); - SimpleCatalog catalog = getSubCatalogForResource(rootCatalog, functionPath); - - if (tvfExists(catalog, functionName)) { - throw new CatalogResourceAlreadyExists(functionFullName); - } + public static void deleteTVFFromCatalog(SimpleCatalog catalog, String fullName) { + if (!tvfExists(catalog, fullName)) { + String errorMessage = String.format("Tried to delete TVF which does not exist: %s", fullName); + throw new CatalogResourceDoesNotExist(fullName, errorMessage); } - } - /** - * Deletes a TVF from the specified paths in a {@link SimpleCatalog} - * - * @param rootCatalog The catalog from which to delete TVFs - * @param functionPaths The paths for the TVF that should be deleted - */ - public static void deleteTVFFromCatalog( - SimpleCatalog rootCatalog, List> functionPaths) { - for (List functionPath : functionPaths) { - String functionName = functionPath.get(functionPath.size() - 1); - SimpleCatalog catalog = getSubCatalogForResource(rootCatalog, functionPath); - - if (tvfExists(catalog, functionName)) { - catalog.removeTableValuedFunction(functionName); - } - } + catalog.removeTableValuedFunction(fullName); } /** - * Creates a TVF in a SimpleCatalog using the provided paths and complying with the provided - * CreateMode. + * Creates a TVF in a {@link SimpleCatalog} using the provided name and complying with the + * provided CreateMode. * - * @param rootCatalog The root SimpleCatalog in which to create the function. - * @param functionPaths The function paths to create the TVF at. If multiple paths are provided, - * multiple copies of the function will be registered in the catalog. - * @param tvfInfo The TVFInfo object representing the TVF that should be created + * @param catalog The catalog in which to create the TVF + * @param nameInCatalog The name under which the TVF will be registered in the catalog + * @param tvfInfo The {@link TVFInfo} object representing the TVF that should be created * @param createMode The CreateMode to use - * @throws CatalogResourceAlreadyExists if the function already exists at any of the provided + * @throws CatalogResourceAlreadyExists if the TVF already exists at any of the provided * paths and CreateMode != CREATE_OR_REPLACE. */ public static void createTVFInCatalog( - SimpleCatalog rootCatalog, - List> functionPaths, + SimpleCatalog catalog, + String nameInCatalog, TVFInfo tvfInfo, CreateMode createMode) { Preconditions.checkArgument( - tvfInfo.getOutputSchema().isPresent(), "Cannot create a a TVF without an output schema"); - - if (createMode.equals(CreateMode.CREATE_OR_REPLACE)) { - deleteTVFFromCatalog(rootCatalog, functionPaths); - } + tvfInfo.getOutputSchema().isPresent(), + "Cannot create a a TVF without an output schema"); + + boolean alreadyExists = tvfExists(catalog, nameInCatalog); + + TableValuedFunction tvf = + new FixedOutputSchemaTVF( + ImmutableList.of(nameInCatalog), + tvfInfo.getSignature(), + tvfInfo.getOutputSchema().get()); + + createResource( + nameInCatalog, + createMode, + "TVF", + alreadyExists, + /*creator=*/() -> catalog.addTableValuedFunction(tvf), + /*deleter=*/() -> deleteTVFFromCatalog(catalog, nameInCatalog) + ); + } - if (createMode.equals(CreateMode.CREATE_DEFAULT)) { - String tvfFullName = String.join(".", tvfInfo.getNamePath()); - validateTVFDoesNotExist(rootCatalog, functionPaths, tvfFullName); + private static void deleteProcedureFromCatalogImpl(SimpleCatalog catalog, String fullName) { + if (!procedureExists(catalog, fullName)) { + String errorMessage = String.format( + "Tried to delete procedure which does not exist: %s", fullName); + throw new CatalogResourceDoesNotExist(fullName, errorMessage); } - for (List functionPath : functionPaths) { - String functionName = functionPath.get(functionPath.size() - 1); - SimpleCatalog catalogForCreation = getSubCatalogForResource(rootCatalog, functionPath); - - TableValuedFunction tvf = - new FixedOutputSchemaTVF( - ImmutableList.of(functionName), - tvfInfo.getSignature(), - tvfInfo.getOutputSchema().get()); - - if (!tvfExists(catalogForCreation, tvf)) { - catalogForCreation.addTableValuedFunction(tvf); - } - } + catalog.removeProcedure(fullName); } /** - * Checks a Procedure does not exist at any of the provided paths. + * Deletes a procedure with the provided name from the {@link SimpleCatalog} * - * @param rootCatalog The catalog to look for procedures in - * @param procedurePaths The list of paths the procedure should not be in - * @param procedureFullName The full name of the procedure we're looking for, only used for error - * reporting - * @throws CatalogResourceAlreadyExists if a procedure exists at any of the provided paths - */ - private static void validateProcedureDoesNotExist( - SimpleCatalog rootCatalog, List> procedurePaths, String procedureFullName) { - for (List procedurePath : procedurePaths) { - String procedureName = procedurePath.get(procedurePath.size() - 1); - SimpleCatalog catalog = getSubCatalogForResource(rootCatalog, procedurePath); - - if (procedureExists(catalog, procedureName)) { - throw new CatalogResourceAlreadyExists(procedureFullName); - } - } - } - - /** - * Deletes a Procedure from the specified paths in a {@link SimpleCatalog} + *

Qualified procedures need to be registered two times in the catalog for analysis to + * work as expected. This method takes care of deleting both copies of the procedure if necessary. * - * @param rootCatalog The catalog from which to delete TVFs - * @param procedurePaths The paths for the Procedure that should be deleted + * @param catalog The catalog from which to delete the procedure + * @param fullName The full name of the procedure in the catalog + * @throws CatalogResourceDoesNotExist if the procedure does not exist in the catalog */ - public static void deleteProcedureFromCatalog( - SimpleCatalog rootCatalog, List> procedurePaths) { - for (List procedurePath : procedurePaths) { - String procedureName = procedurePath.get(procedurePath.size() - 1); - SimpleCatalog catalog = getSubCatalogForResource(rootCatalog, procedurePath); - - if (procedureExists(catalog, procedureName)) { - catalog.removeProcedure(procedureName); - } + public static void deleteProcedureFromCatalog(SimpleCatalog catalog, String fullName) { + deleteProcedureFromCatalogImpl(catalog, fullName); + + if (fullName.contains(".")) { + List nameComponents = Arrays.asList(fullName.split("\\.")); + String nestedName = nameComponents.get(nameComponents.size() - 1); + SimpleCatalog nestedCatalog = getSubCatalogForResource(catalog, nameComponents); + deleteProcedureFromCatalogImpl(nestedCatalog, nestedName); } } /** - * Creates a procedure in a SimpleCatalog using the provided paths and complying with the provided - * CreateMode. + * Creates a procedure in a {@link SimpleCatalog} using the provided name and complying with + * the provided CreateMode. * - * @param rootCatalog The root SimpleCatalog in which to create the procedure. - * @param procedurePaths The procedure paths to create the procedure at. If multiple paths are - * provided, multiple copies of the procedure will be registered in the catalog. + *

Qualified procedures will be registered two times in the catalog for analysis to work as + * expected. A procedure with name "project.dataset.table" will be registered at name paths: + * ["project.dataset.table"] and ["project", "dataset", "table"]. + * + * @param catalog The SimpleCatalog in which to create the procedure + * @param nameInCatalog The name under which the procedure will be registered in the catalog * @param procedureInfo The ProcedureInfo object representing the procedure that should be created * @param createMode The CreateMode to use * @throws CatalogResourceAlreadyExists if the procedure already exists at any of the provided * paths and CreateMode != CREATE_OR_REPLACE. */ public static void createProcedureInCatalog( - SimpleCatalog rootCatalog, - List> procedurePaths, + SimpleCatalog catalog, + String nameInCatalog, ProcedureInfo procedureInfo, CreateMode createMode) { - if (createMode.equals(CreateMode.CREATE_OR_REPLACE)) { - deleteProcedureFromCatalog(rootCatalog, procedurePaths); - } - - if (createMode.equals(CreateMode.CREATE_DEFAULT)) { - String procedureFullName = String.join(".", procedureInfo.getNamePath()); - validateProcedureDoesNotExist(rootCatalog, procedurePaths, procedureFullName); - } - - for (List procedurePath : procedurePaths) { - String procedureName = procedurePath.get(procedurePath.size() - 1); - SimpleCatalog catalogForCreation = getSubCatalogForResource(rootCatalog, procedurePath); + boolean alreadyExists = procedureExists(catalog, nameInCatalog); + Runnable creatorFunction = () -> { Procedure procedure = - new Procedure(ImmutableList.of(procedureName), procedureInfo.getSignature()); - - if (!procedureExists(catalogForCreation, procedure)) { - catalogForCreation.addProcedure(procedure); + new Procedure(ImmutableList.of(nameInCatalog), procedureInfo.getSignature()); + catalog.addProcedure(procedure); + + if (nameInCatalog.contains(".")) { + List nameComponents = Arrays.asList(nameInCatalog.split("\\.")); + String nestedName = nameComponents.get(nameComponents.size() - 1); + SimpleCatalog nestedCatalog = getSubCatalogForResource(catalog, nameComponents); + Procedure nestedProcedure = + new Procedure(ImmutableList.of(nestedName), procedureInfo.getSignature()); + nestedCatalog.addProcedure(nestedProcedure); } - } + }; + + createResource( + nameInCatalog, + createMode, + "Procedure", + alreadyExists, + /*creator=*/ creatorFunction, + /*deleter=*/ () -> deleteProcedureFromCatalog(catalog, nameInCatalog) + ); + } /** @@ -493,11 +413,6 @@ public static void createProcedureInCatalog( * @return The copy of the provided SimpleCatalog. */ public static SimpleCatalog copyCatalog(SimpleCatalog sourceCatalog) { - // Simply serializes and deserializes the source catalog to create a copy. - // This is the most reliable way of creating a copy of a SimpleCatalog, - // as the SimpleCatalog's public interface does not expose enough of the internal - // structures to create an accurate copy. - SimpleCatalogProto serialized = sourceCatalog.serialize(new FileDescriptorSetsBuilder()); - return SimpleCatalog.deserialize(serialized, ImmutableList.of()); + return SimpleCatalogUtil.copyCatalog(sourceCatalog); } } diff --git a/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/basic/BasicCatalogWrapper.java b/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/basic/BasicCatalogWrapper.java index 4e5f660..f04ffe7 100644 --- a/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/basic/BasicCatalogWrapper.java +++ b/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/basic/BasicCatalogWrapper.java @@ -16,11 +16,12 @@ package com.google.zetasql.toolkit.catalog.basic; -import com.google.common.collect.ImmutableList; import com.google.zetasql.Constant; +import com.google.zetasql.LanguageOptions; import com.google.zetasql.SimpleCatalog; import com.google.zetasql.SimpleTable; import com.google.zetasql.ZetaSQLBuiltinFunctionOptions; +import com.google.zetasql.ZetaSQLOptions.LanguageFeature; import com.google.zetasql.resolvedast.ResolvedCreateStatementEnums.CreateMode; import com.google.zetasql.resolvedast.ResolvedCreateStatementEnums.CreateScope; import com.google.zetasql.toolkit.catalog.CatalogOperations; @@ -30,10 +31,8 @@ import com.google.zetasql.toolkit.catalog.TVFInfo; import com.google.zetasql.toolkit.catalog.exceptions.CatalogException; import com.google.zetasql.toolkit.catalog.exceptions.CatalogResourceAlreadyExists; -import java.util.ArrayList; import java.util.Arrays; import java.util.List; -import java.util.stream.Collectors; /** * Basic implementation of CatalogWrapper which does not implement the semantics of any particular @@ -43,9 +42,17 @@ public class BasicCatalogWrapper implements CatalogWrapper { private final SimpleCatalog catalog; + private static final LanguageOptions languageOptionsForFunctionsAndTypes = new LanguageOptions(); + + static { + Arrays.stream(LanguageFeature.values()) + .forEach(languageOptionsForFunctionsAndTypes::enableLanguageFeature); + } + public BasicCatalogWrapper() { this.catalog = new SimpleCatalog("catalog"); - this.catalog.addZetaSQLFunctionsAndTypes(new ZetaSQLBuiltinFunctionOptions()); + this.catalog.addZetaSQLFunctionsAndTypes( + new ZetaSQLBuiltinFunctionOptions(languageOptionsForFunctionsAndTypes)); } public BasicCatalogWrapper(SimpleCatalog initialCatalog) { @@ -59,24 +66,19 @@ public BasicCatalogWrapper(SimpleCatalog initialCatalog) { */ @Override public void register(SimpleTable table, CreateMode createMode, CreateScope createScope) { - CatalogOperations.createTableInCatalog( - this.catalog, - ImmutableList.of(ImmutableList.of(table.getFullName())), - table.getFullName(), - table.getColumnList(), - createMode); + CatalogOperations.createTableInCatalog(this.catalog, table.getFullName(), table, createMode); } @Override public void register(FunctionInfo function, CreateMode createMode, CreateScope createScope) { CatalogOperations.createFunctionInCatalog( - this.catalog, ImmutableList.of(function.getNamePath()), function, createMode); + this.catalog, function.getFullName(), function, createMode); } @Override public void register(TVFInfo tvfInfo, CreateMode createMode, CreateScope createScope) { CatalogOperations.createTVFInCatalog( - this.catalog, ImmutableList.of(tvfInfo.getNamePath()), tvfInfo, createMode); + this.catalog, tvfInfo.getFullName(), tvfInfo, createMode); } /** @@ -87,24 +89,8 @@ public void register(TVFInfo tvfInfo, CreateMode createMode, CreateScope createS @Override public void register( ProcedureInfo procedureInfo, CreateMode createMode, CreateScope createScope) { - if (procedureInfo.getNamePath().size() > 1) { - throw new CatalogException("Procedure name paths should have a single item"); - } - - List namePath = - procedureInfo.getNamePath().stream() - .flatMap(pathElement -> Arrays.stream(pathElement.split("\\."))) - .collect(Collectors.toList()); - String fullName = String.join(".", namePath); - - List> procedurePaths = new ArrayList<>(); - procedurePaths.add(namePath); - if (namePath.size() > 1) { - procedurePaths.add(ImmutableList.of(fullName)); - } - CatalogOperations.createProcedureInCatalog( - this.catalog, procedurePaths, procedureInfo, createMode); + this.catalog, procedureInfo.getFullName(), procedureInfo, createMode); } /** @@ -130,22 +116,22 @@ public void register(Constant constant) { @Override public void removeTable(String fullTableName) { - CatalogOperations.deleteTableFromCatalog(this.catalog, ImmutableList.of(ImmutableList.of(fullTableName))); + CatalogOperations.deleteTableFromCatalog(this.catalog, fullTableName); } @Override public void removeFunction(String function) { - CatalogOperations.deleteFunctionFromCatalog(this.catalog, ImmutableList.of(ImmutableList.of(function))); + CatalogOperations.deleteFunctionFromCatalog(this.catalog, function); } @Override public void removeTVF(String function) { - CatalogOperations.deleteTVFFromCatalog(this.catalog, ImmutableList.of(ImmutableList.of(function))); + CatalogOperations.deleteTVFFromCatalog(this.catalog, function); } @Override public void removeProcedure(String procedure) { - CatalogOperations.deleteProcedureFromCatalog(this.catalog, ImmutableList.of(ImmutableList.of(procedure))); + CatalogOperations.deleteProcedureFromCatalog(this.catalog, procedure); } @Override diff --git a/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/exceptions/CatalogResourceDoesNotExist.java b/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/exceptions/CatalogResourceDoesNotExist.java new file mode 100644 index 0000000..e9f84cf --- /dev/null +++ b/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/exceptions/CatalogResourceDoesNotExist.java @@ -0,0 +1,26 @@ +package com.google.zetasql.toolkit.catalog.exceptions; + +public class CatalogResourceDoesNotExist extends CatalogException { + + private final String resourceName; + + public CatalogResourceDoesNotExist(String resourceName) { + super("Catalog resource already exists: " + resourceName); + this.resourceName = resourceName; + } + + public CatalogResourceDoesNotExist(String resourceName, String message) { + super(message); + this.resourceName = resourceName; + } + + public CatalogResourceDoesNotExist(String resourceName, String message, Throwable cause) { + super(message, cause); + this.resourceName = resourceName; + } + + public String getResourceName() { + return resourceName; + } + +} diff --git a/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/typeparser/ZetaSQLTypeParser.java b/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/typeparser/ZetaSQLTypeParser.java index fa75492..32f408f 100644 --- a/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/typeparser/ZetaSQLTypeParser.java +++ b/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/catalog/typeparser/ZetaSQLTypeParser.java @@ -115,7 +115,7 @@ public Type getResult() { @Override public void exitBasicType(BasicTypeContext ctx) { - String basicTypeName = ctx.BASIC_TYPE().getText(); + String basicTypeName = ctx.BASIC_TYPE().getText().toUpperCase(); TypeKind kind = simpleTypeMapping.getOrDefault(basicTypeName, TypeKind.TYPE_UNKNOWN); Type type = TypeFactory.createSimpleType(kind); this.typeStack.push(type); diff --git a/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/tools/lineage/ColumnLineageExtractor.java b/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/tools/lineage/ColumnLineageExtractor.java index 5e54a2d..4c3f4b3 100644 --- a/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/tools/lineage/ColumnLineageExtractor.java +++ b/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/tools/lineage/ColumnLineageExtractor.java @@ -21,6 +21,7 @@ import com.google.zetasql.Table; import com.google.zetasql.Type; import com.google.zetasql.resolvedast.ResolvedColumn; +import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedQueryStmt; import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedColumnRef; import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedCreateTableAsSelectStmt; import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedCreateViewBase; @@ -49,6 +50,7 @@ * Implements extraction of column-level lineage from ZetaSQL {@link ResolvedStatement}s. * Supported statements: *

    + *
  • SELECT *
  • CREATE TABLE AS SELECT *
  • CREATE [MATERIALIZED] VIEW AS SELECT *
  • INSERT @@ -122,7 +124,7 @@ private static Set extractColumnLevelLineageForOutputColumns( * @param createTableAsSelectStmt The ResolvedCreateTableAsSelectStmt for which to extract lineage * @return The set of resulting {@link ColumnLineage} objects */ - private static Set extractColumnLevelLineage( + public static Set extractColumnLevelLineage( ResolvedCreateTableAsSelectStmt createTableAsSelectStmt) { String fullTableName = String.join(".", createTableAsSelectStmt.getNamePath()); @@ -139,7 +141,7 @@ private static Set extractColumnLevelLineage( * @param createViewBase The ResolvedCreateViewBase statement for which to extract lineage * @return The set of resulting {@link ColumnLineage} objects */ - private static Set extractColumnLevelLineage( + public static Set extractColumnLevelLineage( ResolvedCreateViewBase createViewBase) { String fullViewName = String.join(".", createViewBase.getNamePath()); @@ -149,13 +151,26 @@ private static Set extractColumnLevelLineage( fullViewName, outputColumns, createViewBase); } + /** + * Extracts the column-level lineage entries for a {@link ResolvedQueryStmt} statement + * + * @param statement The ResolvedQueryStmt statement for which to extract lineage + * @param outputTable The name of the table the statement write to + * @return The set of resulting {@link ColumnLineage} objects + */ + public static Set extractColumnLevelLineage( + ResolvedQueryStmt statement, String outputTable) { + List outputColumns = statement.getOutputColumnList(); + return extractColumnLevelLineageForOutputColumns(outputTable, outputColumns, statement); + } + /** * Extracts the column-level lineage entries for a {@link ResolvedInsertStmt} * * @param insertStmt The ResolvedInsertStmt for which to extract lineage * @return The set of resulting {@link ColumnLineage} objects */ - private static Set extractColumnLevelLineage(ResolvedInsertStmt insertStmt) { + public static Set extractColumnLevelLineage(ResolvedInsertStmt insertStmt) { if (Objects.isNull(insertStmt.getQuery())) { // The statement is inserting rows manually using "INSERT INTO ... VALUES ..." // Since it does not query any tables, it does not produce lineage @@ -255,7 +270,7 @@ private static Optional extractColumnLevelLineageForUpdateItem( * @param updateStmt The ResolvedUpdateStmt for which to extract lineage * @return The set of resulting {@link ColumnLineage} objects */ - private static Set extractColumnLevelLineage(ResolvedUpdateStmt updateStmt) { + public static Set extractColumnLevelLineage(ResolvedUpdateStmt updateStmt) { Table targetTable = updateStmt.getTableScan().getTable(); List updateItems = updateStmt.getUpdateItemList(); @@ -317,7 +332,7 @@ private static Set extractColumnLevelLineage( * @param mergeStmt The ResolvedMergeStmt for which to extract lineage * @return The set of resulting {@link ColumnLineage} objects */ - private static Set extractColumnLevelLineage(ResolvedMergeStmt mergeStmt) { + public static Set extractColumnLevelLineage(ResolvedMergeStmt mergeStmt) { Table targetTable = mergeStmt.getTableScan().getTable(); return mergeStmt.getWhenClauseList() @@ -327,34 +342,4 @@ private static Set extractColumnLevelLineage(ResolvedMergeStmt me .collect(Collectors.toSet()); } - /** - * Extracts the column-level lineage entries for a {@link ResolvedStatement}. - * Supported statements: - *
      - *
    • CREATE TABLE AS SELECT - *
    • CREATE [MATERIALIZED] VIEW AS SELECT - *
    • INSERT - *
    • UPDATE - *
    • MERGE - *
    - * - * @param statement The ResolvedStatement for which to extract lineage - * @return The set of resulting {@link ColumnLineage} objects. Empty for unsupported statements. - */ - public static Set extractColumnLevelLineage(ResolvedStatement statement) { - if (statement instanceof ResolvedCreateTableAsSelectStmt) { - return extractColumnLevelLineage((ResolvedCreateTableAsSelectStmt) statement); - } else if (statement instanceof ResolvedInsertStmt) { - return extractColumnLevelLineage((ResolvedInsertStmt) statement); - } else if (statement instanceof ResolvedUpdateStmt) { - return extractColumnLevelLineage((ResolvedUpdateStmt) statement); - } else if (statement instanceof ResolvedMergeStmt) { - return extractColumnLevelLineage((ResolvedMergeStmt) statement); - } else if (statement instanceof ResolvedCreateViewBase) { - return extractColumnLevelLineage((ResolvedCreateViewBase) statement); - } - - return ImmutableSet.of(); - } - } diff --git a/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/tools/lineage/ParentColumnFinder.java b/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/tools/lineage/ParentColumnFinder.java index ee43e46..7410862 100644 --- a/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/tools/lineage/ParentColumnFinder.java +++ b/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/tools/lineage/ParentColumnFinder.java @@ -74,7 +74,7 @@ * those scopes to correlate the ResolvedWithRefScan columns to their parents in the corresponding * ResolvedWithEntry. */ -class ParentColumnFinder extends Visitor { +public class ParentColumnFinder extends Visitor { private final HashMap> columnsToParents = new HashMap<>(); private final Set terminalColumns = new HashSet<>(); diff --git a/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/tools/patch/ZetaSQLPatcher.java b/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/tools/patch/ZetaSQLPatcher.java new file mode 100644 index 0000000..2a0ce55 --- /dev/null +++ b/zetasql-toolkit-core/src/main/java/com/google/zetasql/toolkit/tools/patch/ZetaSQLPatcher.java @@ -0,0 +1,142 @@ +/* + * Copyright 2023 Google LLC All Rights Reserved + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.zetasql.toolkit.tools.patch; + +import com.google.protobuf.ExperimentalApi; +import com.google.protobuf.Message; +import com.google.zetasql.LocalService.AnalyzeRequest; +import com.google.zetasql.LocalService.AnalyzeResponse; +import com.google.zetasql.LocalService.BuildSqlRequest; +import com.google.zetasql.LocalService.BuildSqlResponse; +import com.google.zetasql.LocalService.ParseRequest; +import com.google.zetasql.LocalService.ParseResponse; +import com.google.zetasql.ZetaSqlLocalServiceGrpc; +import com.google.zetasql.io.grpc.MethodDescriptor; +import com.google.zetasql.io.grpc.MethodDescriptor.Marshaller; +import com.google.zetasql.io.grpc.protobuf.ProtoUtils; +import java.lang.reflect.Field; + +/** + * Implements some reflection-based patches to ZetaSQL. The only patch currently available is + * {@link #patchMaxProtobufNestingDepth(int)}. + * + *

    These reflection-based patches are brittle by design and should be avoided whenever possible. + */ +public class ZetaSQLPatcher { + + /** + * Patches the maximum protobuf nesting depth allowed when accessing the ZetaSQL local service + * through GRPC. This patch should be applied when working with large SQL statements that + * hit GRPC's default nesting limit of 100. + * + *

    This new max depth is only set for three key RPCs, which result in a lot of nesting when + * working with large SQL statements. These RPCs are Parse, Analyze and BuildSql. + * + *

    This patch should be considered experimental; as it relies on + * {@link ProtoUtils#marshallerWithRecursionLimit(Message, int)} from grpc-java, which is + * experimental itself. + * + * @param maxDepth the maximum nesting depth to set for RPCs to the ZetaSQL local service. + * Should be greater than 100. + * @throws IllegalAccessException if any reflective access performed by this patch produces an + * illegal access. + */ + @ExperimentalApi + public static void patchMaxProtobufNestingDepth(int maxDepth) + throws IllegalAccessException { + if (!(maxDepth > 100)) { + throw new IllegalArgumentException( + "Invalid max nesting depth for patching protobuf. Should be at least 100, but got: " + + maxDepth); + } + + patchMaxNestingDepthForParse(maxDepth); + patchMaxNestingDepthForAnalyze(maxDepth); + patchMaxNestingDepthForBuildSql(maxDepth); + } + + private static Field getLocalServiceField(String name) { + try { + return ZetaSqlLocalServiceGrpc.class.getDeclaredField(name); + } catch (NoSuchFieldException noSuchFieldException) { + throw new IllegalStateException( + "Tried to access not existent field in class ZetaSqlLocalServiceGrpc: " + name, + noSuchFieldException); + } + } + + private static void patchMaxNestingDepthForAnalyze(int maxDepth) + throws IllegalAccessException { + Field getAnalyzeMethod = getLocalServiceField("getAnalyzeMethod"); + + Marshaller requestMarshaller = ProtoUtils.marshallerWithRecursionLimit( + AnalyzeRequest.getDefaultInstance(), maxDepth); + Marshaller responseMarshaller = ProtoUtils.marshallerWithRecursionLimit( + AnalyzeResponse.getDefaultInstance(), maxDepth); + + MethodDescriptor newMethodDescriptor = + ZetaSqlLocalServiceGrpc.getAnalyzeMethod() + .toBuilder(requestMarshaller, responseMarshaller) + .build(); + + synchronized (ZetaSqlLocalServiceGrpc.class) { + getAnalyzeMethod.setAccessible(true); + getAnalyzeMethod.set(null, newMethodDescriptor); + } + } + + private static void patchMaxNestingDepthForParse(int maxDepth) + throws IllegalAccessException { + Field getParseMethod = getLocalServiceField("getParseMethod"); + + Marshaller requestMarshaller = ProtoUtils.marshallerWithRecursionLimit( + ParseRequest.getDefaultInstance(), maxDepth); + Marshaller responseMarshaller = ProtoUtils.marshallerWithRecursionLimit( + ParseResponse.getDefaultInstance(), maxDepth); + + MethodDescriptor newMethodDescriptor = + ZetaSqlLocalServiceGrpc.getParseMethod() + .toBuilder(requestMarshaller, responseMarshaller) + .build(); + + synchronized (ZetaSqlLocalServiceGrpc.class) { + getParseMethod.setAccessible(true); + getParseMethod.set(null, newMethodDescriptor); + } + } + + private static void patchMaxNestingDepthForBuildSql(int maxDepth) + throws IllegalAccessException { + Field getBuildSqlMethod = getLocalServiceField("getBuildSqlMethod"); + + Marshaller requestMarshaller = ProtoUtils.marshallerWithRecursionLimit( + BuildSqlRequest.getDefaultInstance(), maxDepth); + Marshaller responseMarshaller = ProtoUtils.marshallerWithRecursionLimit( + BuildSqlResponse.getDefaultInstance(), maxDepth); + + MethodDescriptor newMethodDescriptor = + ZetaSqlLocalServiceGrpc.getBuildSqlMethod() + .toBuilder(requestMarshaller, responseMarshaller) + .build(); + + synchronized (ZetaSqlLocalServiceGrpc.class) { + getBuildSqlMethod.setAccessible(true); + getBuildSqlMethod.set(null, newMethodDescriptor); + } + } + +} diff --git a/zetasql-toolkit-core/src/test/java/com/google/zetasql/toolkit/catalog/CatalogOperationsTest.java b/zetasql-toolkit-core/src/test/java/com/google/zetasql/toolkit/catalog/CatalogOperationsTest.java index ece9b01..d00ac40 100644 --- a/zetasql-toolkit-core/src/test/java/com/google/zetasql/toolkit/catalog/CatalogOperationsTest.java +++ b/zetasql-toolkit-core/src/test/java/com/google/zetasql/toolkit/catalog/CatalogOperationsTest.java @@ -41,47 +41,6 @@ private SimpleCatalog createSampleCatalog(String name) { "sample", "column", TypeFactory.createSimpleType(TypeKind.TYPE_STRING)))); catalog.addSimpleTable(sampleTable); - Function function = - new Function( - ImmutableList.of("function"), - "UDF", - ZetaSQLFunctions.FunctionEnums.Mode.SCALAR, - ImmutableList.of( - new FunctionSignature( - new FunctionArgumentType(TypeFactory.createSimpleType(TypeKind.TYPE_STRING)), - ImmutableList.of(), - -1))); - catalog.addFunction(function); - - TVFRelation tvfOutputSchema = - TVFRelation.createColumnBased( - ImmutableList.of( - TVFRelation.Column.create( - "output", TypeFactory.createSimpleType(TypeKind.TYPE_STRING)))); - TableValuedFunction tvf = - new TableValuedFunction.FixedOutputSchemaTVF( - ImmutableList.of("tvf"), - new FunctionSignature( - new FunctionArgumentType( - ZetaSQLFunctions.SignatureArgumentKind.ARG_TYPE_RELATION, - FunctionArgumentType.FunctionArgumentTypeOptions.builder() - .setRelationInputSchema(tvfOutputSchema) - .build(), - 1), - ImmutableList.of(), - -1), - tvfOutputSchema); - catalog.addTableValuedFunction(tvf); - - Procedure procedure = - new Procedure( - ImmutableList.of("procedure"), - new FunctionSignature( - new FunctionArgumentType(ZetaSQLFunctions.SignatureArgumentKind.ARG_TYPE_VOID), - ImmutableList.of(), - -1)); - catalog.addProcedure(procedure); - return catalog; } @@ -112,38 +71,21 @@ void testCreateTableInCatalog() { tableName, "column", TypeFactory.createSimpleType(TypeKind.TYPE_STRING)))); newTable.setFullName(fullTableName); - List newTablePath1 = ImmutableList.of("newTable"); - List newTablePath2 = ImmutableList.of("qualified", "newTable"); - List> newTablePaths = ImmutableList.of(newTablePath1, newTablePath2); - CatalogOperations.createTableInCatalog( - this.testCatalog, - newTablePaths, - fullTableName, - newTable.getColumnList(), - CreateMode.CREATE_DEFAULT); + this.testCatalog, newTable.getFullName(), newTable, CreateMode.CREATE_DEFAULT); - assertAll( - () -> assertTableExists(this.testCatalog, newTablePath1, "Expected created table to exist"), - () -> - assertTableExists(this.testCatalog, newTablePath2, "Expected created table to exist")); + assertTableExists(this.testCatalog, ImmutableList.of("qualified.newTable"), "Expected created table to exist"); } @Test void testDeleteTableFromCatalog() { - List sampleTablePath = ImmutableList.of("sample"); - List nestedSampleTablePath = ImmutableList.of("nested", "sample"); - List> tablePathsToDelete = ImmutableList.of(sampleTablePath, nestedSampleTablePath); - CatalogOperations.deleteTableFromCatalog(this.testCatalog, tablePathsToDelete); + CatalogOperations.deleteTableFromCatalog(this.testCatalog, "sample"); + + assertTableDoesNotExist( + this.testCatalog, ImmutableList.of("sample"), + "Expected table to have been deleted"); - assertAll( - () -> - assertTableDoesNotExist( - this.testCatalog, sampleTablePath, "Expected table to have been deleted"), - () -> - assertTableDoesNotExist( - this.testCatalog, nestedSampleTablePath, "Expected table to have been deleted")); } @Test @@ -156,17 +98,11 @@ void testTableAlreadyExists() { new SimpleColumn( tableName, "column", TypeFactory.createSimpleType(TypeKind.TYPE_INT64)))); - List tablePath = ImmutableList.of("sample"); - assertThrows( CatalogResourceAlreadyExists.class, () -> CatalogOperations.createTableInCatalog( - this.testCatalog, - ImmutableList.of(tablePath), - "sample", - newTable.getColumnList(), - CreateMode.CREATE_DEFAULT)); + this.testCatalog, newTable.getFullName(), newTable, CreateMode.CREATE_DEFAULT)); } @Test @@ -179,17 +115,12 @@ void testReplaceTable() { new SimpleColumn( tableName, "column", TypeFactory.createSimpleType(TypeKind.TYPE_INT64)))); - List tablePath = ImmutableList.of("sample"); CatalogOperations.createTableInCatalog( - this.testCatalog, - ImmutableList.of(tablePath), - "sample", - newTable.getColumnList(), - CreateMode.CREATE_OR_REPLACE); + this.testCatalog, newTable.getFullName(), newTable, CreateMode.CREATE_OR_REPLACE); - Table foundTable = - assertTableExists(this.testCatalog, tablePath, "Expected replaced table to exist"); + Table foundTable = assertTableExists( + this.testCatalog, ImmutableList.of("sample"), "Expected replaced table to exist"); assertEquals( foundTable.getColumn(0).getType(), @@ -212,11 +143,7 @@ void testCreateTableIfNotExists_ExistingTable() throws NotFoundException { Table originalTable = this.testCatalog.findTable(sampleTablePath); CatalogOperations.createTableInCatalog( - this.testCatalog, - ImmutableList.of(sampleTablePath), - "sample", - newTable.getColumnList(), - CreateMode.CREATE_IF_NOT_EXISTS); + this.testCatalog, newTable.getFullName(), newTable, CreateMode.CREATE_IF_NOT_EXISTS); Table foundTable = assertTableExists( @@ -240,16 +167,13 @@ void testCreateTableIfNotExists_NewTable() { new SimpleColumn( tableName, "column", TypeFactory.createSimpleType(TypeKind.TYPE_INT64)))); - List newTablePath = ImmutableList.of("newTable"); - CatalogOperations.createTableInCatalog( - this.testCatalog, - ImmutableList.of(newTablePath), - "newTable", - newTable.getColumnList(), - CreateMode.CREATE_IF_NOT_EXISTS); + this.testCatalog, newTable.getFullName(), newTable, CreateMode.CREATE_IF_NOT_EXISTS); - assertTableExists(this.testCatalog, newTablePath, "Expected table to have been created"); + assertTableExists( + this.testCatalog, + ImmutableList.of("newTable"), + "Expected table to have been created"); } private Function assertFunctionExists(SimpleCatalog catalog, String fullName, String message) { @@ -268,7 +192,7 @@ private void assertFunctionDoesNotExist(SimpleCatalog catalog, String fullName, void testCreateFunctionInCatalog() { FunctionInfo newFunction = FunctionInfo.newBuilder() - .setNamePath(ImmutableList.of("newFunction")) + .setNamePath(ImmutableList.of("qualified.newFunction")) .setGroup("UDF") .setMode(ZetaSQLFunctions.FunctionEnums.Mode.SCALAR) .setSignatures( @@ -280,46 +204,39 @@ void testCreateFunctionInCatalog() { -1))) .build(); - List newFunctionPath1 = ImmutableList.of("newFunction"); - List newFunctionPath2 = ImmutableList.of("qualified", "newFunction"); - List> newFunctionPaths = ImmutableList.of(newFunctionPath1, newFunctionPath2); - CatalogOperations.createFunctionInCatalog( - this.testCatalog, newFunctionPaths, newFunction, CreateMode.CREATE_DEFAULT); - - SimpleCatalog qualifiedNestedCatalog = this.testCatalog.getCatalog("qualified", null); + this.testCatalog, "qualified.newFunction", newFunction, CreateMode.CREATE_DEFAULT); - assertNotNull( - qualifiedNestedCatalog, - "Expected the nested catalog to exist after creating a resource in it"); + assertFunctionExists( + this.testCatalog, "UDF:qualified.newFunction", "Expected created function to exist"); - assertAll( - () -> - assertFunctionExists( - this.testCatalog, "UDF:newFunction", "Expected created function to exist"), - () -> - assertFunctionExists( - qualifiedNestedCatalog, "UDF:newFunction", "Expected created function to exist")); } @Test void testDeleteFunctionFromCatalog() { - List sampleFunctionPath = ImmutableList.of("function"); - List nestedSampleFunctionPath = ImmutableList.of("nested", "function"); + FunctionInfo newFunction = + FunctionInfo.newBuilder() + .setNamePath(ImmutableList.of("qualified.newFunction")) + .setGroup("UDF") + .setMode(ZetaSQLFunctions.FunctionEnums.Mode.SCALAR) + .setSignatures( + ImmutableList.of( + new FunctionSignature( + new FunctionArgumentType( + TypeFactory.createSimpleType(TypeKind.TYPE_STRING)), + ImmutableList.of(), + -1))) + .build(); - List> functionPathsToDelete = - ImmutableList.of(sampleFunctionPath, nestedSampleFunctionPath); - CatalogOperations.deleteFunctionFromCatalog(this.testCatalog, functionPathsToDelete); + CatalogOperations.createFunctionInCatalog( + this.testCatalog, "qualified.newFunction", newFunction, CreateMode.CREATE_DEFAULT); + CatalogOperations.deleteFunctionFromCatalog(this.testCatalog, "qualified.newFunction"); + + assertFunctionDoesNotExist( + this.testCatalog, + "UDF:qualified.newFunction", + "Expected function to have been deleted"); - assertAll( - () -> - assertFunctionDoesNotExist( - this.testCatalog.getCatalog("nested", null), - "UDF:function", - "Expected function to have been deleted"), - () -> - assertFunctionDoesNotExist( - this.testCatalog, "UDF:function", "Expected function to have been deleted")); } private TableValuedFunction assertTVFExists(SimpleCatalog catalog, String name, String message) { @@ -338,7 +255,7 @@ private void assertTVFDoesNotExist(SimpleCatalog catalog, String name, String me void testCreateTVFInCatalog() { TVFInfo newTVF = TVFInfo.newBuilder() - .setNamePath(ImmutableList.of("newTVF")) + .setNamePath(ImmutableList.of("qualified.newTVF")) .setSignature( new FunctionSignature( new FunctionArgumentType( @@ -350,44 +267,33 @@ void testCreateTVFInCatalog() { TypeFactory.createSimpleType(TypeKind.TYPE_STRING))) .build(); - List newFunctionPath1 = ImmutableList.of("newTVF"); - List newFunctionPath2 = ImmutableList.of("qualified", "newTVF"); - List> newFunctionPaths = ImmutableList.of(newFunctionPath1, newFunctionPath2); - CatalogOperations.createTVFInCatalog( - this.testCatalog, newFunctionPaths, newTVF, CreateMode.CREATE_DEFAULT); - - SimpleCatalog qualifiedNestedCatalog = this.testCatalog.getCatalog("qualified", null); + this.testCatalog, "qualified.newTVF", newTVF, CreateMode.CREATE_DEFAULT); - assertNotNull( - qualifiedNestedCatalog, - "Expected the nested catalog to exist after creating a resource in it"); - - assertAll( - () -> assertTVFExists(this.testCatalog, "newTVF", "Expected created function to exist"), - () -> - assertTVFExists( - qualifiedNestedCatalog, "newTVF", "Expected created function to exist")); + assertTVFExists(testCatalog, "qualified.newTVF", "Expected created function to exist"); } @Test void testDeleteTVFFromCatalog() { - List sampleFunctionPath = ImmutableList.of("tvf"); - List nestedSampleFunctionPath = ImmutableList.of("nested", "tvf"); + TVFInfo tvf = TVFInfo.newBuilder() + .setNamePath(ImmutableList.of("qualified.newTVF")) + .setSignature( + new FunctionSignature( + new FunctionArgumentType( + ZetaSQLFunctions.SignatureArgumentKind.ARG_TYPE_RELATION), + ImmutableList.of(), + -1)) + .setOutputSchema( + TVFRelation.createValueTableBased( + TypeFactory.createSimpleType(TypeKind.TYPE_STRING))) + .build(); - List> functionPathsToDelete = - ImmutableList.of(sampleFunctionPath, nestedSampleFunctionPath); - CatalogOperations.deleteTVFFromCatalog(this.testCatalog, functionPathsToDelete); + CatalogOperations.createTVFInCatalog( + this.testCatalog, "qualified.newTVF", tvf, CreateMode.CREATE_DEFAULT); + CatalogOperations.deleteTVFFromCatalog(this.testCatalog, "qualified.newTVF"); - assertAll( - () -> - assertTVFDoesNotExist( - this.testCatalog.getCatalog("nested", null), - "tvf", - "Expected function to have been deleted"), - () -> - assertTVFDoesNotExist( - this.testCatalog, "tvf", "Expected function to have been deleted")); + assertTVFDoesNotExist( + this.testCatalog, "qualified.newTVF", "Expected function to have been deleted"); } private Procedure assertProcedureExists( @@ -404,18 +310,17 @@ private void assertProcedureDoesNotExist( void testCreateProcedureInCatalog() { ProcedureInfo newProcedure = new ProcedureInfo( - ImmutableList.of("newProcedure"), + ImmutableList.of("qualified.newProcedure"), new FunctionSignature( new FunctionArgumentType(ZetaSQLFunctions.SignatureArgumentKind.ARG_TYPE_VOID), ImmutableList.of(), -1)); - List newProcedurePath1 = ImmutableList.of("newProcedure"); + List newProcedurePath1 = ImmutableList.of("qualified.newProcedure"); List newProcedurePath2 = ImmutableList.of("qualified", "newProcedure"); - List> newProcedurePaths = ImmutableList.of(newProcedurePath1, newProcedurePath2); CatalogOperations.createProcedureInCatalog( - this.testCatalog, newProcedurePaths, newProcedure, CreateMode.CREATE_DEFAULT); + this.testCatalog, "qualified.newProcedure", newProcedure, CreateMode.CREATE_DEFAULT); assertAll( () -> @@ -428,20 +333,30 @@ void testCreateProcedureInCatalog() { @Test void testDeleteProcedureFromCatalog() { - List sampleProcedurePath = ImmutableList.of("procedure"); - List nestedSampleProcedurePath = ImmutableList.of("nested", "procedure"); + ProcedureInfo newProcedure = + new ProcedureInfo( + ImmutableList.of("qualified.newProcedure"), + new FunctionSignature( + new FunctionArgumentType(ZetaSQLFunctions.SignatureArgumentKind.ARG_TYPE_VOID), + ImmutableList.of(), + -1)); + + CatalogOperations.createProcedureInCatalog( + this.testCatalog, "qualified.newProcedure", newProcedure, CreateMode.CREATE_DEFAULT); + + List procedurePath1 = ImmutableList.of("newProcedure"); + List procedurePath2 = ImmutableList.of("qualified", "newProcedure"); - List> pathsToDelete = ImmutableList.of(sampleProcedurePath, nestedSampleProcedurePath); - CatalogOperations.deleteProcedureFromCatalog(this.testCatalog, pathsToDelete); + CatalogOperations.deleteProcedureFromCatalog(this.testCatalog, "qualified.newProcedure"); assertAll( () -> assertProcedureDoesNotExist( - this.testCatalog, sampleProcedurePath, "Expected procedure to have been deleted"), + this.testCatalog, procedurePath1, "Expected procedure to have been deleted"), () -> assertProcedureDoesNotExist( this.testCatalog, - nestedSampleProcedurePath, + procedurePath2, "Expected procedure to have been deleted")); } diff --git a/zetasql-toolkit-core/src/test/java/com/google/zetasql/toolkit/catalog/typeparser/ZetaSQLTypeParserTest.java b/zetasql-toolkit-core/src/test/java/com/google/zetasql/toolkit/catalog/typeparser/ZetaSQLTypeParserTest.java index 88de699..577bcfb 100644 --- a/zetasql-toolkit-core/src/test/java/com/google/zetasql/toolkit/catalog/typeparser/ZetaSQLTypeParserTest.java +++ b/zetasql-toolkit-core/src/test/java/com/google/zetasql/toolkit/catalog/typeparser/ZetaSQLTypeParserTest.java @@ -36,6 +36,7 @@ void parseSimpleTypes() { Map inputsToExpectedKinds = new HashMap<>(); inputsToExpectedKinds.put("STRING", TypeKind.TYPE_STRING); inputsToExpectedKinds.put("INT64", TypeKind.TYPE_INT64); + inputsToExpectedKinds.put("DATE", TypeKind.TYPE_DATE); inputsToExpectedKinds.put("NUMERIC", TypeKind.TYPE_NUMERIC); inputsToExpectedKinds.put("INTERVAL", TypeKind.TYPE_INTERVAL); inputsToExpectedKinds.put("JSON", TypeKind.TYPE_JSON); @@ -46,8 +47,31 @@ void parseSimpleTypes() { inputToExpectedKind -> () -> assertEquals( + TypeFactory.createSimpleType(inputToExpectedKind.getValue()), ZetaSQLTypeParser.parse(inputToExpectedKind.getKey()), + "Failed to parse type: " + inputToExpectedKind.getKey())); + + assertAll(assertions); + } + + @Test + void typesAreCaseInsensitive() { + Map inputsToExpectedKinds = new HashMap<>(); + inputsToExpectedKinds.put("string", TypeKind.TYPE_STRING); + inputsToExpectedKinds.put("inT64", TypeKind.TYPE_INT64); + inputsToExpectedKinds.put("date", TypeKind.TYPE_DATE); + inputsToExpectedKinds.put("NuMEric", TypeKind.TYPE_NUMERIC); + inputsToExpectedKinds.put("intErval", TypeKind.TYPE_INTERVAL); + inputsToExpectedKinds.put("JSOn", TypeKind.TYPE_JSON); + + Stream assertions = + inputsToExpectedKinds.entrySet().stream() + .map( + inputToExpectedKind -> + () -> + assertEquals( TypeFactory.createSimpleType(inputToExpectedKind.getValue()), + ZetaSQLTypeParser.parse(inputToExpectedKind.getKey()), "Failed to parse type: " + inputToExpectedKind.getKey())); assertAll(assertions); @@ -58,13 +82,13 @@ void parseSimpleTypesWithParameters() { assertAll( () -> assertEquals( - ZetaSQLTypeParser.parse("STRING(MAX)"), TypeFactory.createSimpleType(TypeKind.TYPE_STRING), + ZetaSQLTypeParser.parse("STRING(MAX)"), "Failed to parse string type with parameter: STRING(10)"), () -> assertEquals( - ZetaSQLTypeParser.parse("NUMERIC(10, 2)"), TypeFactory.createSimpleType(TypeKind.TYPE_NUMERIC), + ZetaSQLTypeParser.parse("NUMERIC(10, 2)"), "Failed to parse numeric type with parameters: NUMERIC(10, 2)")); } @@ -75,7 +99,7 @@ void parseArray() { TypeFactory.createArrayType(TypeFactory.createSimpleType(TypeKind.TYPE_STRING)); assertEquals( - ZetaSQLTypeParser.parse(typeStr), expectedType, "Failed to parse type ARRAY"); + expectedType, ZetaSQLTypeParser.parse(typeStr), "Failed to parse type ARRAY"); } @Test @@ -88,8 +112,8 @@ void parseStruct() { new StructField("f2", TypeFactory.createSimpleType(TypeKind.TYPE_INT64)))); assertEquals( - ZetaSQLTypeParser.parse(typeStr), expectedType, + ZetaSQLTypeParser.parse(typeStr), "Failed to parse type STRUCT"); } @@ -103,8 +127,8 @@ void parseStructWithParameterType() { new StructField("f2", TypeFactory.createSimpleType(TypeKind.TYPE_NUMERIC)))); assertEquals( - ZetaSQLTypeParser.parse(typeStr), expectedType, + ZetaSQLTypeParser.parse(typeStr), "Failed to parse struct type STRUCT"); } @@ -119,8 +143,8 @@ void parseArrayOfStructs() { Type expectedType = TypeFactory.createArrayType(structType); assertEquals( - ZetaSQLTypeParser.parse(typeStr), expectedType, + ZetaSQLTypeParser.parse(typeStr), "Failed to array of structs ARRAY>"); } @@ -139,8 +163,8 @@ void parseStructWithMultipleNestingLevels() { ImmutableList.of(new StructField("f1", TypeFactory.createArrayType(innerStructType)))); assertEquals( - ZetaSQLTypeParser.parse(typeStr), expectedType, + ZetaSQLTypeParser.parse(typeStr), "Failed to struct with multiple nesting levels " + "STRUCT, f1_2 NUMERIC(10, 2)>>>"); } diff --git a/zetasql-toolkit-core/src/test/java/com/google/zetasql/toolkit/tools/patch/ZetaSQLPatcherTest.java b/zetasql-toolkit-core/src/test/java/com/google/zetasql/toolkit/tools/patch/ZetaSQLPatcherTest.java new file mode 100644 index 0000000..01fff8c --- /dev/null +++ b/zetasql-toolkit-core/src/test/java/com/google/zetasql/toolkit/tools/patch/ZetaSQLPatcherTest.java @@ -0,0 +1,68 @@ +/* + * Copyright 2023 Google LLC All Rights Reserved + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.zetasql.toolkit.tools.patch; + +import static org.junit.jupiter.api.Assertions.*; + +import com.google.zetasql.AnalyzerOptions; +import com.google.zetasql.Parser; +import com.google.zetasql.toolkit.ZetaSQLToolkitAnalyzer; +import org.junit.jupiter.api.Assumptions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class ZetaSQLPatcherTest { + + private ZetaSQLToolkitAnalyzer analyzer; + + @BeforeEach + void init() { + AnalyzerOptions analyzerOptions = new AnalyzerOptions(); + analyzerOptions.getLanguageOptions().enableMaximumLanguageFeatures(); + analyzerOptions.getLanguageOptions().setSupportsAllStatementKinds(); + + this.analyzer = new ZetaSQLToolkitAnalyzer(analyzerOptions); + } + + public static String generateNestedSelectStatement(int times) { + if (times == 1) { + return "SELECT 1"; + } + + return String.format("SELECT 1 FROM (%s)", generateNestedSelectStatement(times - 1)); + } + @Test + void testMaxNestingDepthPatch() { + // The query is a SELECT statement nested 100 times. Parsing or analyzing + // such a statement normally fails due to reaching the default max nesting + // depth in the ZetaSQL local service. + // After patching the max nesting depth, it should not fail. + String query = generateNestedSelectStatement(100); + + try { + ZetaSQLPatcher.patchMaxProtobufNestingDepth(1000); + } catch (IllegalAccessException err) { + Assumptions.abort("Aborting test because the patch was not applied successfully due to " + + "disallowed reflective access."); + } + + assertDoesNotThrow(() -> { + this.analyzer.analyzeStatements(query).next(); + }); + } + +} diff --git a/zetasql-toolkit-examples/pom.xml b/zetasql-toolkit-examples/pom.xml index 3cad956..0c9433c 100644 --- a/zetasql-toolkit-examples/pom.xml +++ b/zetasql-toolkit-examples/pom.xml @@ -22,12 +22,12 @@ com.google.zetasql.toolkit zetasql-toolkit - 0.4.1 + 0.5.0 ../pom.xml zetasql-toolkit-examples - 0.4.1 + 0.5.0 ${project.groupId}:${project.artifactId} diff --git a/zetasql-toolkit-examples/src/main/java/com/google/zetasql/toolkit/examples/ExtractColumnLevelLineage.java b/zetasql-toolkit-examples/src/main/java/com/google/zetasql/toolkit/examples/ExtractColumnLevelLineage.java index 5ac74c5..43355b7 100644 --- a/zetasql-toolkit-examples/src/main/java/com/google/zetasql/toolkit/examples/ExtractColumnLevelLineage.java +++ b/zetasql-toolkit-examples/src/main/java/com/google/zetasql/toolkit/examples/ExtractColumnLevelLineage.java @@ -18,7 +18,11 @@ import com.google.common.collect.ImmutableList; import com.google.zetasql.AnalyzerOptions; +import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedCreateTableAsSelectStmt; +import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedInsertStmt; +import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedMergeStmt; import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedStatement; +import com.google.zetasql.resolvedast.ResolvedNodes.ResolvedUpdateStmt; import com.google.zetasql.toolkit.AnalyzedStatement; import com.google.zetasql.toolkit.ZetaSQLToolkitAnalyzer; import com.google.zetasql.toolkit.catalog.bigquery.BigQueryCatalog; @@ -60,8 +64,11 @@ private static void lineageForCreateTableAsSelectStatement( Iterator statementIterator = analyzer.analyzeStatements(query, catalog); ResolvedStatement statement = statementIterator.next().getResolvedStatement().get(); + ResolvedCreateTableAsSelectStmt createTableAsSelectStmt = + (ResolvedCreateTableAsSelectStmt) statement; - Set lineageEntries = ColumnLineageExtractor.extractColumnLevelLineage(statement); + Set lineageEntries = + ColumnLineageExtractor.extractColumnLevelLineage(createTableAsSelectStmt); System.out.println("Extracted column lineage from CREATE TABLE AS SELECT"); outputLineage(query, lineageEntries); @@ -84,8 +91,10 @@ private static void lineageForInsertStatement( Iterator statementIterator = analyzer.analyzeStatements(query, catalog); ResolvedStatement statement = statementIterator.next().getResolvedStatement().get(); + ResolvedInsertStmt insertStmt = (ResolvedInsertStmt) statement; - Set lineageEntries = ColumnLineageExtractor.extractColumnLevelLineage(statement); + Set lineageEntries = + ColumnLineageExtractor.extractColumnLevelLineage(insertStmt); System.out.println("Extracted column lineage from INSERT"); outputLineage(query, lineageEntries); @@ -101,8 +110,10 @@ private static void lineageForUpdateStatement( Iterator statementIterator = analyzer.analyzeStatements(query, catalog); ResolvedStatement statement = statementIterator.next().getResolvedStatement().get(); + ResolvedUpdateStmt updateStmt = (ResolvedUpdateStmt) statement; - Set lineageEntries = ColumnLineageExtractor.extractColumnLevelLineage(statement); + Set lineageEntries = + ColumnLineageExtractor.extractColumnLevelLineage(updateStmt); System.out.println("Extracted column lineage from UPDATE"); outputLineage(query, lineageEntries); @@ -121,8 +132,9 @@ private static void lineageForMergeStatement( Iterator statementIterator = analyzer.analyzeStatements(query, catalog); ResolvedStatement statement = statementIterator.next().getResolvedStatement().get(); + ResolvedMergeStmt mergeStmt = (ResolvedMergeStmt) statement; - Set lineageEntries = ColumnLineageExtractor.extractColumnLevelLineage(statement); + Set lineageEntries = ColumnLineageExtractor.extractColumnLevelLineage(mergeStmt); System.out.println("Extracted column lineage from MERGE"); outputLineage(query, lineageEntries); diff --git a/zetasql-toolkit-spanner/pom.xml b/zetasql-toolkit-spanner/pom.xml index 3238a71..40f73f6 100644 --- a/zetasql-toolkit-spanner/pom.xml +++ b/zetasql-toolkit-spanner/pom.xml @@ -6,12 +6,12 @@ com.google.zetasql.toolkit zetasql-toolkit - 0.4.1 + 0.5.0 ../pom.xml zetasql-toolkit-spanner - 0.4.1 + 0.5.0 ${project.groupId}:${project.artifactId} diff --git a/zetasql-toolkit-spanner/src/main/java/com/google/zetasql/toolkit/catalog/spanner/SpannerCatalog.java b/zetasql-toolkit-spanner/src/main/java/com/google/zetasql/toolkit/catalog/spanner/SpannerCatalog.java index 27b7d6b..27edc5d 100644 --- a/zetasql-toolkit-spanner/src/main/java/com/google/zetasql/toolkit/catalog/spanner/SpannerCatalog.java +++ b/zetasql-toolkit-spanner/src/main/java/com/google/zetasql/toolkit/catalog/spanner/SpannerCatalog.java @@ -107,6 +107,7 @@ private SpannerCatalog( * @param projectId The Spanner project id * @param instance The Spanner instance name * @param database The Spanner database name + * @return the new SpannerCatalog instance */ public static SpannerCatalog usingSpannerClient( String projectId, String instance, String database) { @@ -123,6 +124,7 @@ public static SpannerCatalog usingSpannerClient( * @param instance The Spanner instance name * @param database The Spanner database name * @param spannerClient The Spanner client to use + * @return the new SpannerCatalog instance */ public static SpannerCatalog usingSpannerClient( String projectId, String instance, String database, Spanner spannerClient) { @@ -136,6 +138,7 @@ public static SpannerCatalog usingSpannerClient( * {@link CatalogResources} object. * * @param resources The {@link CatalogResources} object from which this catalog will get tables + * @return the new SpannerCatalog instance */ public static SpannerCatalog usingResources(CatalogResources resources) { SpannerResourceProvider resourceProvider = new LocalSpannerResourceProvider(resources); @@ -151,18 +154,13 @@ public static SpannerCatalog usingResources(CatalogResources resources) { */ @Override public void register(SimpleTable table, CreateMode createMode, CreateScope createScope) { - String tableName = table.getName(); + String fullName = table.getFullName(); - if (tableName.contains(".")) { - throw new InvalidSpannerTableName(tableName); + if (fullName.contains(".")) { + throw new InvalidSpannerTableName(fullName); } - CatalogOperations.createTableInCatalog( - this.catalog, - ImmutableList.of(ImmutableList.of(table.getName())), - table.getName(), - table.getColumnList(), - createMode); + CatalogOperations.createTableInCatalog(this.catalog, table.getFullName(), table, createMode); } @Override @@ -191,7 +189,7 @@ public void register(Constant constant) { @Override public void removeTable(String table) { this.validateSpannerTableNames(ImmutableList.of(table)); - CatalogOperations.deleteTableFromCatalog(this.catalog, ImmutableList.of(ImmutableList.of(table))); + CatalogOperations.deleteTableFromCatalog(this.catalog, table); } @Override