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