From 082ce917159d3f789b973e0e2de5ed0440e4dca7 Mon Sep 17 00:00:00 2001 From: Ahmed Abualsaud Date: Tue, 11 Aug 2026 14:27:26 -0700 Subject: [PATCH 1/2] call procedure --- .../sql/src/main/codegen/config.fmpp | 3 + .../src/main/codegen/includes/parserImpls.ftl | 32 ++ .../sql/impl/parser/SqlCallProcedure.java | 363 ++++++++++++++ .../extensions/sql/meta/catalog/Catalog.java | 9 + .../sql/meta/catalog/Procedure.java | 55 +++ .../sql/impl/parser/SqlCallProcedureTest.java | 458 ++++++++++++++++++ 6 files changed, 920 insertions(+) create mode 100644 sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlCallProcedure.java create mode 100644 sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/catalog/Procedure.java create mode 100644 sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlCallProcedureTest.java diff --git a/sdks/java/extensions/sql/src/main/codegen/config.fmpp b/sdks/java/extensions/sql/src/main/codegen/config.fmpp index c3692430610f..473ec1942e84 100644 --- a/sdks/java/extensions/sql/src/main/codegen/config.fmpp +++ b/sdks/java/extensions/sql/src/main/codegen/config.fmpp @@ -38,6 +38,7 @@ data: { "org.apache.beam.sdk.extensions.sql.impl.parser.SqlSetOptionBeam" "org.apache.beam.sdk.extensions.sql.impl.parser.SqlAlterCatalog" "org.apache.beam.sdk.extensions.sql.impl.parser.SqlAlterTable" + "org.apache.beam.sdk.extensions.sql.impl.parser.SqlCallProcedure" "org.apache.beam.sdk.extensions.sql.impl.utils.CalciteUtils" "org.apache.beam.sdk.schemas.Schema" ] @@ -414,6 +415,7 @@ data: { # List of non-reserved keywords to add; # items in this list become non-reserved nonReservedKeywordsToAdd: [ + "SYSTEM" ] # List of non-reserved keywords to remove; @@ -439,6 +441,7 @@ data: { "SqlSetOptionBeam(Span.of(), null)" "SqlAlterCatalog(Span.of(), null)" "SqlAlterTable(Span.of(), null)" + "SqlCallProcedure(Span.of())" ] # List of methods for parsing custom literals. diff --git a/sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl b/sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl index cb8eec438728..c310cbeff472 100644 --- a/sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl +++ b/sdks/java/extensions/sql/src/main/codegen/includes/parserImpls.ftl @@ -855,6 +855,38 @@ Schema.FieldType SimpleType() : } } +/** + * CALL ( catalog_name '.' )? ( 'system' '.' )? procedure_name + * '(' ( arg ( ',' arg )* )? ')' + * + * where arg := literal | param_name '=>' literal + */ +SqlCall SqlCallProcedure(Span s) : +{ + final SqlIdentifier procedureName; + final List args = new ArrayList(); +} +{ + { + s.add(this); + } + procedureName = CompoundIdentifier() + + [ + AddArg0(args, ExprContext.ACCEPT_NONCURSOR) + ( + { + checkNonQueryExpression(ExprContext.ACCEPT_NONCURSOR); + } + AddArg(args, ExprContext.ACCEPT_NONCURSOR) + )* + ] + + { + return new SqlCallProcedure(s.end(this), procedureName, args); + } +} + SqlSetOptionBeam SqlSetOptionBeam(Span s, String scope) : { SqlIdentifier name; diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlCallProcedure.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlCallProcedure.java new file mode 100644 index 000000000000..8ae5ea6effa5 --- /dev/null +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlCallProcedure.java @@ -0,0 +1,363 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.beam.sdk.extensions.sql.impl.parser; + +import static org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.util.Static.RESOURCE; + +import java.math.BigDecimal; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import org.apache.beam.sdk.extensions.sql.impl.CatalogManagerSchema; +import org.apache.beam.sdk.extensions.sql.meta.catalog.Catalog; +import org.apache.beam.sdk.extensions.sql.meta.catalog.Procedure; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.jdbc.CalcitePrepare; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.jdbc.CalciteSchema; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlCall; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlIdentifier; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlKind; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlLiteral; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlNode; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlNumericLiteral; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlOperator; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlSpecialOperator; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlUtil; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlWriter; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.parser.SqlParserPos; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.type.SqlTypeName; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.util.Pair; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.checkerframework.checker.nullness.qual.Nullable; + +/** + * A {@code CALL} statement invoking a stored {@link Procedure}: + * + * 
{@code
+ * CALL [catalog_name.][system.]procedure_name(arg1, arg2, ...)
+ * CALL [catalog_name.][system.]procedure_name(param2 => arg2, param1 => arg1, ...)
+ * }
+ * + *

Arguments are passed either by position or by name, the namespace component (if any) must + * be {@code system}, and procedure names + * resolve case-insensitively. Procedures are provided by the target {@link Catalog} via {@link + * Catalog#loadProcedure(String)}. + */ +public class SqlCallProcedure extends SqlCall implements BeamSqlParser.ExecutableStatement { + private static final SqlOperator OPERATOR = new SqlSpecialOperator("CALL", SqlKind.OTHER_DDL); + private static final String SYSTEM_NAMESPACE = "system"; + + private final SqlIdentifier procedureName; + private final List args; + + public SqlCallProcedure(SqlParserPos pos, SqlIdentifier procedureName, List args) { + super(pos); + this.procedureName = procedureName; + this.args = ImmutableList.copyOf(args); + } + + @Override + public SqlOperator getOperator() { + return OPERATOR; + } + + @Override + public List getOperandList() { + return ImmutableList.builder().add(procedureName).addAll(args).build(); + } + + @Override + public void unparse(SqlWriter writer, int leftPrec, int rightPrec) { + writer.keyword("CALL"); + procedureName.unparse(writer, leftPrec, rightPrec); + SqlWriter.Frame frame = writer.startList(SqlWriter.FrameTypeEnum.FUN_CALL, "(", ")"); + for (SqlNode arg : args) { + writer.sep(","); + if (arg.getKind() == SqlKind.ARGUMENT_ASSIGNMENT) { + SqlCall assignment = (SqlCall) arg; + assignment.operand(1).unparse(writer, 0, 0); + writer.keyword("=>"); + assignment.operand(0).unparse(writer, 0, 0); + } else { + arg.unparse(writer, 0, 0); + } + } + writer.endList(frame); + } + + @Override + public void execute(CalcitePrepare.Context context) { + List path = procedureName.names; + @Nullable String catalogName = null; + String procName; + switch (path.size()) { + case 1: + procName = path.get(0); + break; + case 2: + checkSystemNamespace(path.get(0)); + procName = path.get(1); + break; + case 3: + catalogName = path.get(0); + checkSystemNamespace(path.get(1)); + procName = path.get(2); + break; + default: + throw invalidProcedureNameException(); + } + + Catalog catalog = resolveCatalog(context, catalogName, procName); + + // Procedure resolution is case-insensitive + String lookupName = procName.toLowerCase(Locale.ROOT); + @Nullable Procedure procedure = catalog.loadProcedure(lookupName); + if (procedure == null) { + throw SqlUtil.newContextException( + procedureName.getParserPosition(), + RESOURCE.internal( + String.format( + "Procedure '%s' not found in catalog '%s' (type '%s').", + lookupName, catalog.name(), catalog.type()))); + } + + procedure.execute(bindArguments(procedure)); + } + + /** Resolves the target {@link Catalog}, defaulting to the currently active one. */ + private Catalog resolveCatalog( + CalcitePrepare.Context context, @Nullable String catalogName, String procName) { + final Pair pair = + SqlDdlNodes.schema( + context, true, new SqlIdentifier(procName, procedureName.getParserPosition())); + org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.schema.Schema schema = + pair.left.schema; + if (!(schema instanceof CatalogManagerSchema)) { + throw SqlUtil.newContextException( + procedureName.getParserPosition(), + RESOURCE.internal( + "Attempting to execute 'CALL' with unexpected Calcite Schema of type " + + schema.getClass())); + } + CatalogManagerSchema catalogManagerSchema = (CatalogManagerSchema) schema; + + if (catalogName == null) { + return catalogManagerSchema.getCurrentCatalogSchema().getCatalog(); + } + for (Catalog catalog : catalogManagerSchema.catalogs()) { + if (catalog.name().equals(catalogName)) { + return catalog; + } + } + throw SqlUtil.newContextException( + procedureName.getParserPosition(), + RESOURCE.internal(String.format("Catalog '%s' not found.", catalogName))); + } + + private void checkSystemNamespace(String namespace) { + if (!namespace.equalsIgnoreCase(SYSTEM_NAMESPACE)) { + throw invalidProcedureNameException(); + } + } + + private RuntimeException invalidProcedureNameException() { + return SqlUtil.newContextException( + procedureName.getParserPosition(), + RESOURCE.internal( + String.format( + "Invalid procedure name '%s': expected 'procedure_name', " + + "'system.procedure_name', or 'catalog.system.procedure_name'.", + String.join(".", procedureName.names)))); + } + + /** + * Validates the arguments against the procedure's declared parameters and binds them to a {@link + * Row} over {@link Procedure#parameters()}. + */ + private Row bindArguments(Procedure procedure) { + Schema parameters = procedure.parameters(); + String procName = procedure.name(); + int paramCount = parameters.getFieldCount(); + + boolean hasNamed = args.stream().anyMatch(a -> a.getKind() == SqlKind.ARGUMENT_ASSIGNMENT); + boolean hasPositional = args.stream().anyMatch(a -> a.getKind() != SqlKind.ARGUMENT_ASSIGNMENT); + if (hasNamed && hasPositional) { + throw SqlUtil.newContextException( + procedureName.getParserPosition(), + RESOURCE.internal("Mixing named and positional arguments is not supported.")); + } + + // Map of parameter name -> argument value node. + Map providedArgs = new HashMap<>(); + if (hasNamed) { + for (SqlNode arg : args) { + SqlCall assignment = (SqlCall) arg; + SqlNode value = assignment.operand(0); + SqlIdentifier nameId = assignment.operand(1); + // Parameter names resolve case-insensitively; parameters are declared in lower_snake_case. + String name = nameId.getSimple().toLowerCase(Locale.ROOT); + if (!parameters.hasField(name)) { + throw SqlUtil.newContextException( + nameId.getParserPosition(), + RESOURCE.internal( + String.format( + "Procedure '%s' does not accept an argument named '%s'. " + + "Expected parameters: %s", + procName, nameId.getSimple(), parameters.getFieldNames()))); + } + if (providedArgs.put(name, value) != null) { + throw SqlUtil.newContextException( + nameId.getParserPosition(), + RESOURCE.internal( + String.format( + "Duplicate argument name '%s' in call to procedure '%s'.", name, procName))); + } + } + } else { + if (args.size() > paramCount) { + throw SqlUtil.newContextException( + procedureName.getParserPosition(), + RESOURCE.internal( + String.format( + "Too many arguments for procedure '%s': expected at most %s, got %s.", + procName, paramCount, args.size()))); + } + for (int i = 0; i < args.size(); i++) { + providedArgs.put(parameters.getField(i).getName(), args.get(i)); + } + } + + List missingRequired = new ArrayList<>(); + @Nullable Object[] values = new @Nullable Object[paramCount]; + for (int i = 0; i < paramCount; i++) { + Schema.Field field = parameters.getField(i); + @Nullable SqlNode valueNode = providedArgs.get(field.getName()); + if (valueNode == null) { + if (!field.getType().getNullable()) { + missingRequired.add(field.getName()); + } + continue; + } + @Nullable Object value = convertLiteral(procName, field, valueNode); + if (value == null && !field.getType().getNullable()) { + throw SqlUtil.newContextException( + valueNode.getParserPosition(), + RESOURCE.internal( + String.format( + "Argument '%s' of procedure '%s' is required and cannot be NULL.", + field.getName(), procName))); + } + values[i] = value; + } + if (!missingRequired.isEmpty()) { + throw SqlUtil.newContextException( + procedureName.getParserPosition(), + RESOURCE.internal( + String.format( + "Missing required argument(s) for procedure '%s': %s.", + procName, missingRequired))); + } + + return Row.withSchema(parameters).addValues(Arrays.asList(values)).build(); + } + + /** Converts a literal argument node to a Java value of the parameter's declared type. */ + private @Nullable Object convertLiteral(String procName, Schema.Field field, SqlNode node) { + if (!(node instanceof SqlLiteral)) { + throw SqlUtil.newContextException( + node.getParserPosition(), + RESOURCE.internal( + String.format( + "Argument '%s' of procedure '%s' must be a literal value.", + field.getName(), procName))); + } + SqlLiteral literal = (SqlLiteral) node; + if (literal.getTypeName() == SqlTypeName.NULL) { + return null; + } + + Schema.TypeName typeName = field.getType().getTypeName(); + try { + switch (typeName) { + case STRING: + checkLiteralType(procName, field, literal, literal.getTypeName() == SqlTypeName.CHAR); + return literal.getValueAs(String.class); + case BOOLEAN: + checkLiteralType(procName, field, literal, literal.getTypeName() == SqlTypeName.BOOLEAN); + return literal.getValueAs(Boolean.class); + case BYTE: + return exactNumeric(procName, field, literal).byteValueExact(); + case INT16: + return exactNumeric(procName, field, literal).shortValueExact(); + case INT32: + return exactNumeric(procName, field, literal).intValueExact(); + case INT64: + return exactNumeric(procName, field, literal).longValueExact(); + case FLOAT: + checkLiteralType(procName, field, literal, literal instanceof SqlNumericLiteral); + return literal.getValueAs(BigDecimal.class).floatValue(); + case DOUBLE: + checkLiteralType(procName, field, literal, literal instanceof SqlNumericLiteral); + return literal.getValueAs(BigDecimal.class).doubleValue(); + case DECIMAL: + checkLiteralType(procName, field, literal, literal instanceof SqlNumericLiteral); + return literal.getValueAs(BigDecimal.class); + default: + throw SqlUtil.newContextException( + node.getParserPosition(), + RESOURCE.internal( + String.format( + "Parameter '%s' of procedure '%s' has type %s, which is not yet supported " + + "for CALL arguments.", + field.getName(), procName, typeName))); + } + } catch (ArithmeticException e) { + throw typeMismatchException(procName, field, literal); + } + } + + private BigDecimal exactNumeric(String procName, Schema.Field field, SqlLiteral literal) { + checkLiteralType( + procName, + field, + literal, + literal instanceof SqlNumericLiteral && ((SqlNumericLiteral) literal).isExact()); + return literal.getValueAs(BigDecimal.class); + } + + private void checkLiteralType( + String procName, Schema.Field field, SqlLiteral literal, boolean typeMatches) { + if (!typeMatches) { + throw typeMismatchException(procName, field, literal); + } + } + + private RuntimeException typeMismatchException( + String procName, Schema.Field field, SqlLiteral literal) { + return SqlUtil.newContextException( + literal.getParserPosition(), + RESOURCE.internal( + String.format( + "Argument '%s' of procedure '%s' expects type %s, but got: %s", + field.getName(), procName, field.getType().getTypeName(), literal))); + } +} diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/catalog/Catalog.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/catalog/Catalog.java index cbf1b45c31e7..2eab466725a1 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/catalog/Catalog.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/catalog/Catalog.java @@ -98,4 +98,13 @@ public interface Catalog { * Returns all the {@link TableProvider}s available to this {@link Catalog}, organized by type. */ Map tableProviders(); + + /** + * Returns the {@link Procedure} registered under this name, or null if this catalog does not + * provide it. Procedure resolution is case-insensitive: implementations always receive a + * lowercase name. + */ + default @Nullable Procedure loadProcedure(String name) { + return null; + } } diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/catalog/Procedure.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/catalog/Procedure.java new file mode 100644 index 000000000000..b720aa56c11a --- /dev/null +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/meta/catalog/Procedure.java @@ -0,0 +1,55 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.beam.sdk.extensions.sql.meta.catalog; + +import org.apache.beam.sdk.annotations.Internal; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.values.Row; + +/** + * A stored procedure that can be invoked with the SQL {@code CALL} statement: + * + * 

{@code
+ * CALL [catalog_name.][system.]procedure_name(arg1, arg2, ...);
+ * CALL [catalog_name.][system.]procedure_name(param2 => arg2, param1 => arg1, ...);
+ * }
+ * + *

Procedures are provided by a {@link Catalog} via {@link Catalog#loadProcedure(String)}. + * Arguments may be passed by position or by name (but not both in the same call), and are validated + * and bound against {@link #parameters()} before {@link #execute(Row)} is invoked. + */ +@Internal +public interface Procedure { + + /** The name of this procedure, in {@code lower_snake_case} (e.g. {@code "add_files"}). */ + String name(); + + /** + * Declares this procedure's parameters. + * + *

Field order defines the positional-argument order. Fields that are non-nullable are required + * arguments; nullable fields are optional and default to null when omitted. + */ + Schema parameters(); + + /** + * Runs the procedure. {@code args} is a {@link Row} over {@link #parameters()} holding the bound + * argument values, with omitted optional arguments set to null. + */ + void execute(Row args); +} diff --git a/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlCallProcedureTest.java b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlCallProcedureTest.java new file mode 100644 index 000000000000..a65256b1e933 --- /dev/null +++ b/sdks/java/extensions/sql/src/test/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlCallProcedureTest.java @@ -0,0 +1,458 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.beam.sdk.extensions.sql.impl.parser; + +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.containsString; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import com.google.auto.service.AutoService; +import java.util.Arrays; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Supplier; +import org.apache.beam.sdk.extensions.sql.BeamSqlCli; +import org.apache.beam.sdk.extensions.sql.impl.BeamSqlEnv; +import org.apache.beam.sdk.extensions.sql.impl.ParseException; +import org.apache.beam.sdk.extensions.sql.meta.catalog.Catalog; +import org.apache.beam.sdk.extensions.sql.meta.catalog.CatalogRegistrar; +import org.apache.beam.sdk.extensions.sql.meta.catalog.InMemoryCatalog; +import org.apache.beam.sdk.extensions.sql.meta.catalog.InMemoryCatalogManager; +import org.apache.beam.sdk.extensions.sql.meta.catalog.Procedure; +import org.apache.beam.sdk.options.PipelineOptionsFactory; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.runtime.CalciteContextException; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlIdentifier; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.SqlLiteral; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.dialect.AnsiSqlDialect; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.fun.SqlStdOperatorTable; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.parser.SqlParserPos; +import org.apache.beam.vendor.calcite.v1_40_0.org.apache.calcite.sql.pretty.SqlPrettyWriter; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.junit.Before; +import org.junit.Test; + +/** + * Unit tests for {@code CALL} procedure statements ({@link SqlCallProcedure}). + * + *

Uses a test-only {@link Catalog} that exposes a recording test procedure. + */ +public class SqlCallProcedureTest { + private static final String TEST_CATALOG_TYPE = "test_with_procedures"; + + private InMemoryCatalogManager catalogManager; + private BeamSqlEnv env; + + @Before + public void setUp() { + TestProcedure.reset(); + NoOpProcedure.reset(); + catalogManager = new InMemoryCatalogManager(); + env = + BeamSqlEnv.builder(catalogManager) + .setPipelineOptions(PipelineOptionsFactory.create()) + .build(); + env.executeDdl("CREATE CATALOG test_cat TYPE '" + TEST_CATALOG_TYPE + "'"); + } + + private void useTestCatalog() { + env.executeDdl("USE CATALOG test_cat"); + } + + @Test + public void testCallProcedure_positionalArgs() { + useTestCatalog(); + env.executeDdl("CALL test_proc('db.tbl', 12345, true, 0.75, 'hello')"); + + assertEquals(1, TestProcedure.executeCount.get()); + assertEquals( + Row.withSchema(TestProcedure.PARAMETERS) + .addValues("db.tbl", 12345L, true, 0.75, "hello") + .build(), + TestProcedure.lastArgs.get()); + } + + @Test + public void testCallProcedure_positionalArgs_omitsTrailingOptional() { + useTestCatalog(); + env.executeDdl("CALL test_proc('db.tbl', 42)"); + + assertEquals(1, TestProcedure.executeCount.get()); + assertEquals( + Row.withSchema(TestProcedure.PARAMETERS).addValues("db.tbl", 42L, null, null, null).build(), + TestProcedure.lastArgs.get()); + } + + @Test + public void testCallProcedure_namedArgs_anyOrder() { + useTestCatalog(); + env.executeDdl("CALL test_proc(snapshot_id => 42, note => 'n', target => 'db.tbl')"); + + assertEquals(1, TestProcedure.executeCount.get()); + assertEquals( + Row.withSchema(TestProcedure.PARAMETERS).addValues("db.tbl", 42L, null, null, "n").build(), + TestProcedure.lastArgs.get()); + } + + @Test + public void testCallProcedure_emptyArgs() { + useTestCatalog(); + env.executeDdl("CALL noop_proc()"); + + assertEquals(1, NoOpProcedure.executeCount.get()); + } + + @Test + public void testCallProcedure_negativeNumberArg() { + useTestCatalog(); + env.executeDdl("CALL test_proc('db.tbl', -42)"); + + Row args = TestProcedure.lastArgs.get(); + assertEquals(Long.valueOf(-42L), args.getInt64("snapshot_id")); + assertEquals(1, TestProcedure.executeCount.get()); + } + + @Test + public void testCallProcedure_nullForOptionalArg() { + useTestCatalog(); + env.executeDdl("CALL test_proc('db.tbl', 1, NULL)"); + + Row args = TestProcedure.lastArgs.get(); + assertNull(args.getBoolean("use_caching")); + assertEquals(1, TestProcedure.executeCount.get()); + } + + @Test + public void testCallProcedure_mixedArgs_error() { + useTestCatalog(); + CalciteContextException e = + assertThrows( + CalciteContextException.class, + () -> env.executeDdl("CALL test_proc('db.tbl', snapshot_id => 42)")); + assertThat( + e.getMessage(), containsString("Mixing named and positional arguments is not supported")); + assertEquals(0, TestProcedure.executeCount.get()); + } + + @Test + public void testCallProcedure_missingRequiredArg_error() { + useTestCatalog(); + CalciteContextException e = + assertThrows( + CalciteContextException.class, + () -> env.executeDdl("CALL test_proc(target => 'db.tbl')")); + assertThat(e.getMessage(), containsString("Missing required argument(s)")); + assertThat(e.getMessage(), containsString("snapshot_id")); + } + + @Test + public void testCallProcedure_unknownArgName_error() { + useTestCatalog(); + CalciteContextException e = + assertThrows( + CalciteContextException.class, + () -> + env.executeDdl( + "CALL test_proc(target => 'db.tbl', snapshot_id => 1, bad_param => 2)")); + assertThat(e.getMessage(), containsString("does not accept an argument named 'bad_param'")); + } + + @Test + public void testCallProcedure_duplicateArgName_error() { + useTestCatalog(); + CalciteContextException e = + assertThrows( + CalciteContextException.class, + () -> env.executeDdl("CALL test_proc(target => 'a', target => 'b', snapshot_id => 1)")); + assertThat(e.getMessage(), containsString("Duplicate argument name 'target'")); + } + + @Test + public void testCallProcedure_tooManyArgs_error() { + useTestCatalog(); + CalciteContextException e = + assertThrows( + CalciteContextException.class, + () -> env.executeDdl("CALL test_proc('a', 1, true, 0.5, 'n', 'extra')")); + assertThat(e.getMessage(), containsString("Too many arguments")); + assertThat(e.getMessage(), containsString("expected at most 5, got 6")); + } + + @Test + public void testCallProcedure_nonLiteralArg_error() { + useTestCatalog(); + CalciteContextException e = + assertThrows( + CalciteContextException.class, () -> env.executeDdl("CALL test_proc('db.tbl', 1 + 2)")); + assertThat(e.getMessage(), containsString("must be a literal")); + } + + @Test + public void testCallProcedure_typeMismatch_error() { + useTestCatalog(); + CalciteContextException e = + assertThrows( + CalciteContextException.class, + () -> env.executeDdl("CALL test_proc('db.tbl', 'not_a_number')")); + assertThat(e.getMessage(), containsString("snapshot_id")); + assertThat(e.getMessage(), containsString("INT64")); + } + + @Test + public void testCallProcedure_nullForRequiredArg_error() { + useTestCatalog(); + CalciteContextException e = + assertThrows( + CalciteContextException.class, + () -> env.executeDdl("CALL test_proc(target => NULL, snapshot_id => 1)")); + assertThat(e.getMessage(), containsString("'target'")); + assertThat(e.getMessage(), containsString("cannot be NULL")); + } + + @Test + public void testCallProcedure_systemNamespace() { + useTestCatalog(); + env.executeDdl("CALL system.test_proc('db.tbl', 7)"); + + Row args = TestProcedure.lastArgs.get(); + assertEquals(Long.valueOf(7L), args.getInt64("snapshot_id")); + assertEquals(1, TestProcedure.executeCount.get()); + } + + @Test + public void testCallProcedure_fullyQualified() { + // Current catalog remains 'default'; qualify the test catalog explicitly. + env.executeDdl("CALL test_cat.system.test_proc('db.tbl', 99)"); + + Row args = TestProcedure.lastArgs.get(); + assertEquals(Long.valueOf(99L), args.getInt64("snapshot_id")); + assertEquals(1, TestProcedure.executeCount.get()); + } + + @Test + public void testCallProcedure_caseInsensitiveResolution() { + useTestCatalog(); + env.executeDdl("CALL SYSTEM.TEST_PROC('db.tbl', 3)"); + + assertEquals(1, TestProcedure.executeCount.get()); + assertEquals(1, TestProcedure.executeCount.get()); + } + + @Test + public void testCallProcedure_twoPartNonSystemNamespace_error() { + useTestCatalog(); + CalciteContextException e = + assertThrows( + CalciteContextException.class, () -> env.executeDdl("CALL foo.test_proc('db.tbl', 1)")); + assertThat(e.getMessage(), containsString("Invalid procedure name 'foo.test_proc'")); + } + + @Test + public void testCallProcedure_threePartNonSystemNamespace_error() { + CalciteContextException e = + assertThrows( + CalciteContextException.class, + () -> env.executeDdl("CALL test_cat.foo.test_proc('db.tbl', 1)")); + assertThat(e.getMessage(), containsString("Invalid procedure name 'test_cat.foo.test_proc'")); + } + + @Test + public void testCallProcedure_tooManyNameParts_error() { + CalciteContextException e = + assertThrows( + CalciteContextException.class, () -> env.executeDdl("CALL a.b.c.d('db.tbl', 1)")); + assertThat(e.getMessage(), containsString("Invalid procedure name 'a.b.c.d'")); + } + + @Test + public void testCallProcedure_unknownCatalog_error() { + CalciteContextException e = + assertThrows( + CalciteContextException.class, + () -> env.executeDdl("CALL nope.system.test_proc('db.tbl', 1)")); + assertThat(e.getMessage(), containsString("Catalog 'nope' not found")); + } + + @Test + public void testCallProcedure_unknownProcedure_error() { + useTestCatalog(); + CalciteContextException e = + assertThrows(CalciteContextException.class, () -> env.executeDdl("CALL nope_proc()")); + assertThat(e.getMessage(), containsString("Procedure 'nope_proc' not found in catalog")); + assertThat(e.getMessage(), containsString("test_cat")); + } + + @Test + public void testCallProcedure_catalogWithoutProcedureSupport_error() { + // The default in-memory catalog provides no procedures. + CalciteContextException e = + assertThrows( + CalciteContextException.class, () -> env.executeDdl("CALL test_proc('db.tbl', 1)")); + assertThat(e.getMessage(), containsString("Procedure 'test_proc' not found in catalog")); + assertThat(e.getMessage(), containsString("default")); + } + + @Test + public void testCallProcedure_missingParens_error() { + useTestCatalog(); + assertThrows(ParseException.class, () -> env.executeDdl("CALL test_proc")); + } + + @Test + public void testCallProcedure_isDdl() { + assertTrue(env.isDdl("CALL test_proc('db.tbl', 1)")); + } + + @Test + public void testCallProcedure_throughBeamSqlCli() { + InMemoryCatalogManager cliCatalogManager = new InMemoryCatalogManager(); + BeamSqlCli cli = new BeamSqlCli().catalogManager(cliCatalogManager); + cli.execute("CREATE CATALOG cli_cat TYPE '" + TEST_CATALOG_TYPE + "'"); + cli.execute("USE CATALOG cli_cat"); + cli.execute("CALL test_proc('db.tbl', 5)"); + + assertEquals(1, TestProcedure.executeCount.get()); + Row args = TestProcedure.lastArgs.get(); + assertEquals(Long.valueOf(5L), args.getInt64("snapshot_id")); + } + + @Test + public void testUnparseCallProcedure() { + SqlCallProcedure call = + new SqlCallProcedure( + SqlParserPos.ZERO, + new SqlIdentifier(Arrays.asList("my_cat", "system", "test_proc"), SqlParserPos.ZERO), + ImmutableList.of( + SqlLiteral.createCharString("db.tbl", SqlParserPos.ZERO), + SqlStdOperatorTable.ARGUMENT_ASSIGNMENT.createCall( + SqlParserPos.ZERO, + SqlLiteral.createExactNumeric("5", SqlParserPos.ZERO), + new SqlIdentifier("snapshot_id", SqlParserPos.ZERO)))); + + SqlPrettyWriter sqlWriter = + new SqlPrettyWriter(SqlPrettyWriter.config().withDialect(AnsiSqlDialect.DEFAULT)); + call.unparse(sqlWriter, 0, 0); + assertEquals( + "CALL `my_cat`.`system`.`test_proc`('db.tbl', `snapshot_id` => 5)", + sqlWriter.toSqlString().getSql()); + } + + /** + * Test-only catalog exposing procedures: a static registry map of procedure names to suppliers. + */ + public static class TestProcedureCatalog extends InMemoryCatalog { + private static final Map> PROCEDURES = + ImmutableMap.of( + TestProcedure.NAME, TestProcedure::new, + NoOpProcedure.NAME, NoOpProcedure::new); + + public TestProcedureCatalog(String name, Map properties) { + super(name, properties); + } + + @Override + public String type() { + return TEST_CATALOG_TYPE; + } + + @Override + public @Nullable Procedure loadProcedure(String name) { + Supplier supplier = PROCEDURES.get(name); + return supplier == null ? null : supplier.get(); + } + } + + /** Registers {@link TestProcedureCatalog} for {@code CREATE CATALOG ... TYPE}. */ + @AutoService(CatalogRegistrar.class) + public static class TestProcedureCatalogRegistrar implements CatalogRegistrar { + @Override + public Iterable> getCatalogs() { + return ImmutableList.of(TestProcedureCatalog.class); + } + } + + /** A recording test procedure with required and optional parameters. */ + private static class TestProcedure implements Procedure { + static final String NAME = "test_proc"; + static final Schema PARAMETERS = + Schema.builder() + .addStringField("target") + .addInt64Field("snapshot_id") + .addNullableBooleanField("use_caching") + .addNullableDoubleField("min_ratio") + .addNullableStringField("note") + .build(); + + private static final AtomicReference<@Nullable Row> lastArgs = new AtomicReference<>(); + private static final AtomicInteger executeCount = new AtomicInteger(); + + static void reset() { + lastArgs.set(null); + executeCount.set(0); + } + + @Override + public String name() { + return NAME; + } + + @Override + public Schema parameters() { + return PARAMETERS; + } + + @Override + public void execute(Row args) { + lastArgs.set(args); + executeCount.incrementAndGet(); + } + } + + /** A recording test procedure with no parameters. */ + private static class NoOpProcedure implements Procedure { + static final String NAME = "noop_proc"; + + private static final AtomicInteger executeCount = new AtomicInteger(); + + static void reset() { + executeCount.set(0); + } + + @Override + public String name() { + return NAME; + } + + @Override + public Schema parameters() { + return Schema.builder().build(); + } + + @Override + public void execute(Row args) { + executeCount.incrementAndGet(); + } + } +} From bb8f52c6bf63859f0f245705191cc17db591a9f3 Mon Sep 17 00:00:00 2001 From: Ahmed Abualsaud Date: Tue, 11 Aug 2026 16:06:35 -0700 Subject: [PATCH 2/2] spotless --- .../sdk/extensions/sql/impl/parser/SqlCallProcedure.java | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlCallProcedure.java b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlCallProcedure.java index 8ae5ea6effa5..e5194c37da36 100644 --- a/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlCallProcedure.java +++ b/sdks/java/extensions/sql/src/main/java/org/apache/beam/sdk/extensions/sql/impl/parser/SqlCallProcedure.java @@ -57,10 +57,9 @@ * CALL [catalog_name.][system.]procedure_name(param2 => arg2, param1 => arg1, ...) * } * - *

Arguments are passed either by position or by name, the namespace component (if any) must - * be {@code system}, and procedure names - * resolve case-insensitively. Procedures are provided by the target {@link Catalog} via {@link - * Catalog#loadProcedure(String)}. + *

Arguments are passed either by position or by name, the namespace component (if any) must be + * {@code system}, and procedure names resolve case-insensitively. Procedures are provided by the + * target {@link Catalog} via {@link Catalog#loadProcedure(String)}. */ public class SqlCallProcedure extends SqlCall implements BeamSqlParser.ExecutableStatement { private static final SqlOperator OPERATOR = new SqlSpecialOperator("CALL", SqlKind.OTHER_DDL);