Examples with JdbcConnection org.apache.beam.sdk.extensions.sql.impl.JdbcConnection used on opensource projects

Example 1 with JdbcConnection

use of org.apache.beam.sdk.extensions.sql.impl.JdbcConnection in project beam by apache.

the class BeamZetaSqlCatalogTest method rejectsCreateFunctionStmtWithUnsupportedReturnType.

@Test
public void rejectsCreateFunctionStmtWithUnsupportedReturnType() {
    JdbcConnection jdbcConnection = createJdbcConnection();
    AnalyzerOptions analyzerOptions = SqlAnalyzer.baseAnalyzerOptions();
    BeamZetaSqlCatalog beamCatalog = BeamZetaSqlCatalog.create(jdbcConnection.getCurrentSchemaPlus(), jdbcConnection.getTypeFactory(), analyzerOptions);
    String sql = "CREATE FUNCTION foo() RETURNS ARRAY<TIME> LANGUAGE java OPTIONS (path='/does/not/exist');";
    ResolvedNodes.ResolvedStatement resolvedStatement = Analyzer.analyzeStatement(sql, analyzerOptions, beamCatalog.getZetaSqlCatalog());
    ResolvedNodes.ResolvedCreateFunctionStmt createFunctionStmt = (ResolvedNodes.ResolvedCreateFunctionStmt) resolvedStatement;
    thrown.expect(UnsupportedOperationException.class);
    thrown.expectMessage("ZetaSQL type TYPE_TIME not allowed in function foo");
    beamCatalog.addFunction(createFunctionStmt);
}

Example 2 with JdbcConnection

use of org.apache.beam.sdk.extensions.sql.impl.JdbcConnection in project beam by apache.

the class BeamZetaSqlCatalogTest method rejectsScalarFunctionImplWithUnsupportedParameterType.

@Test
public void rejectsScalarFunctionImplWithUnsupportedParameterType() throws NoSuchMethodException {
    JdbcConnection jdbcConnection = createJdbcConnection();
    SchemaPlus calciteSchema = jdbcConnection.getCurrentSchemaPlus();
    Method method = TakesArrayTimeFn.class.getMethod("eval", List.class);
    calciteSchema.add("take_array", ScalarFunctionImpl.create(method));
    thrown.expect(UnsupportedOperationException.class);
    thrown.expectMessage("Calcite type TIME not allowed in function take_array");
    BeamZetaSqlCatalog.create(calciteSchema, jdbcConnection.getTypeFactory(), SqlAnalyzer.baseAnalyzerOptions());
}

Example 3 with JdbcConnection

use of org.apache.beam.sdk.extensions.sql.impl.JdbcConnection in project beam by apache.

the class BeamZetaSqlCatalogTest method rejectsCreateFunctionStmtWithUnsupportedParameterType.

@Test
public void rejectsCreateFunctionStmtWithUnsupportedParameterType() {
    JdbcConnection jdbcConnection = createJdbcConnection();
    AnalyzerOptions analyzerOptions = SqlAnalyzer.baseAnalyzerOptions();
    BeamZetaSqlCatalog beamCatalog = BeamZetaSqlCatalog.create(jdbcConnection.getCurrentSchemaPlus(), jdbcConnection.getTypeFactory(), analyzerOptions);
    String sql = "CREATE FUNCTION foo(a ARRAY<TIME>) RETURNS INT64 LANGUAGE java OPTIONS (path='/does/not/exist');";
    ResolvedNodes.ResolvedStatement resolvedStatement = Analyzer.analyzeStatement(sql, analyzerOptions, beamCatalog.getZetaSqlCatalog());
    ResolvedNodes.ResolvedCreateFunctionStmt createFunctionStmt = (ResolvedNodes.ResolvedCreateFunctionStmt) resolvedStatement;
    thrown.expect(UnsupportedOperationException.class);
    thrown.expectMessage("ZetaSQL type TYPE_TIME not allowed in function foo");
    beamCatalog.addFunction(createFunctionStmt);
}

Example 4 with JdbcConnection

use of org.apache.beam.sdk.extensions.sql.impl.JdbcConnection in project beam by apache.

the class ZetaSqlJavaUdfTest method testUdfFromCatalog.

/**
 * This tests a subset of the code path used by {@link #testSqlTransformRegisterUdf()}.
 */
@Test
public void testUdfFromCatalog() throws NoSuchMethodException {
    // Add IncrementFn to Calcite schema.
    JdbcConnection jdbcConnection = JdbcDriver.connect(new ReadOnlyTableProvider("empty_table_provider", ImmutableMap.of()), PipelineOptionsFactory.create());
    Method method = IncrementFn.class.getMethod("eval", Long.class);
    jdbcConnection.getCurrentSchemaPlus().add("increment", ScalarFunctionImpl.create(method));
    this.config = Frameworks.newConfigBuilder(config).defaultSchema(jdbcConnection.getCurrentSchemaPlus()).build();
    String sql = "SELECT increment(0);";
    ZetaSQLQueryPlanner zetaSQLQueryPlanner = new ZetaSQLQueryPlanner(config);
    BeamRelNode beamRelNode = zetaSQLQueryPlanner.convertToBeamRel(sql);
    PCollection<Row> stream = BeamSqlRelUtils.toPCollection(pipeline, beamRelNode);
    final Schema schema = Schema.builder().addInt64Field("field1").build();
    PAssert.that(stream).containsInAnyOrder(Row.withSchema(schema).addValues(1L).build());
    pipeline.run().waitUntilFinish(Duration.standardMinutes(PIPELINE_EXECUTION_WAITTIME_MINUTES));
}

Example 5 with JdbcConnection

use of org.apache.beam.sdk.extensions.sql.impl.JdbcConnection in project beam by apache.

the class ZetaSqlJavaUdfTypeTest method setUp.

@Before
public void setUp() throws NoSuchMethodException {
    initialize();
    // Register test table.
    JdbcConnection jdbcConnection = JdbcDriver.connect(new ReadOnlyTableProvider("table_provider", ImmutableMap.of("table", table)), PipelineOptionsFactory.create());
    // Register UDFs.
    SchemaPlus schema = jdbcConnection.getCurrentSchemaPlus();
    schema.add("test_boolean", ScalarFunctionImpl.create(BooleanIdentityFn.class.getMethod("eval", Boolean.class)));
    schema.add("test_int64", ScalarFunctionImpl.create(Int64IdentityFn.class.getMethod("eval", Long.class)));
    schema.add("test_string", ScalarFunctionImpl.create(StringIdentityFn.class.getMethod("eval", String.class)));
    schema.add("test_bytes", ScalarFunctionImpl.create(BytesIdentityFn.class.getMethod("eval", byte[].class)));
    schema.add("test_float64", ScalarFunctionImpl.create(DoubleIdentityFn.class.getMethod("eval", Double.class)));
    schema.add("test_date", ScalarFunctionImpl.create(DateIdentityFn.class.getMethod("eval", Date.class)));
    schema.add("test_timestamp", ScalarFunctionImpl.create(TimestampIdentityFn.class.getMethod("eval", Timestamp.class)));
    schema.add("test_array", ScalarFunctionImpl.create(ListIdentityFn.class.getMethod("eval", List.class)));
    schema.add("test_numeric", ScalarFunctionImpl.create(BigDecimalIdentityFn.class.getMethod("eval", BigDecimal.class)));
    this.config = Frameworks.newConfigBuilder(config).defaultSchema(schema).build();
}