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

Example 6 with JdbcConnection

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

the class ZetaSqlTestBase method initialize.

protected void initialize() {
    JdbcConnection jdbcConnection = JdbcDriver.connect(createBeamTableProvider(), PipelineOptionsFactory.create());
    this.config = Frameworks.newConfigBuilder().defaultSchema(jdbcConnection.getCurrentSchemaPlus()).traitDefs(ImmutableList.of(ConventionTraitDef.INSTANCE)).context(Contexts.of(jdbcConnection.config())).ruleSets(ZetaSQLQueryPlanner.getZetaSqlRuleSets().toArray(new RuleSet[0])).costFactory(BeamCostModel.FACTORY).typeSystem(jdbcConnection.getTypeFactory().getTypeSystem()).build();
}

Example 7 with JdbcConnection

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

the class PubsubTableProviderIT method connect.

@SuppressWarnings("unchecked")
private CalciteConnection connect(PipelineOptions options, TableProvider... tableProviders) {
    // HACK: PipelineOptions should expose a prominent method to do this reliably
    // The actual options are in the "options" field of the converted map
    Map<String, String> argsMap = ((Map<String, Object>) MAPPER.convertValue(pipeline.getOptions(), Map.class).get("options")).entrySet().stream().filter((entry) -> {
        if (entry.getValue() instanceof List) {
            if (!((List) entry.getValue()).isEmpty()) {
                throw new IllegalArgumentException("Cannot encode list arguments");
            }
            // We can encode empty lists, just omit them.
            return false;
        }
        return true;
    }).collect(Collectors.toMap(Map.Entry::getKey, entry -> toArg(entry.getValue())));
    InMemoryMetaStore inMemoryMetaStore = new InMemoryMetaStore();
    for (TableProvider tableProvider : tableProviders) {
        inMemoryMetaStore.registerProvider(tableProvider);
    }
    JdbcConnection connection = JdbcDriver.connect(inMemoryMetaStore, options);
    connection.setPipelineOptionsMap(argsMap);
    return connection;
}

Example 8 with JdbcConnection

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

the class ZetaSQLPushDownTest method initializeCalciteEnvironmentWithContext.

private static void initializeCalciteEnvironmentWithContext(Context... extraContext) {
    JdbcConnection jdbcConnection = JdbcDriver.connect(tableProvider, PipelineOptionsFactory.create());
    SchemaPlus defaultSchemaPlus = jdbcConnection.getCurrentSchemaPlus();
    final ImmutableList<RelTraitDef> traitDefs = ImmutableList.of(ConventionTraitDef.INSTANCE);
    Object[] contexts = ImmutableList.<Context>builder().add(Contexts.of(jdbcConnection.config())).add(extraContext).build().toArray();
    config = Frameworks.newConfigBuilder().defaultSchema(defaultSchemaPlus).traitDefs(traitDefs).context(Contexts.of(contexts)).ruleSets(ZetaSQLQueryPlanner.getZetaSqlRuleSets().toArray(new RuleSet[0])).costFactory(BeamCostModel.FACTORY).typeSystem(jdbcConnection.getTypeFactory().getTypeSystem()).build();
}

Example 9 with JdbcConnection

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

the class BeamZetaSqlCatalogTest method loadsUserDefinedFunctionsFromSchema.

@Test
public void loadsUserDefinedFunctionsFromSchema() throws NoSuchMethodException {
    JdbcConnection jdbcConnection = createJdbcConnection();
    SchemaPlus calciteSchema = jdbcConnection.getCurrentSchemaPlus();
    Method method = IncrementFn.class.getMethod("eval", Long.class);
    calciteSchema.add("increment", ScalarFunctionImpl.create(method));
    BeamZetaSqlCatalog beamCatalog = BeamZetaSqlCatalog.create(calciteSchema, jdbcConnection.getTypeFactory(), SqlAnalyzer.baseAnalyzerOptions());
    assertNotNull("ZetaSQL catalog contains function signature.", beamCatalog.getZetaSqlCatalog().getFunctionByFullName(USER_DEFINED_JAVA_SCALAR_FUNCTIONS + ":increment"));
    assertEquals("Beam catalog contains function definition.", UserFunctionDefinitions.JavaScalarFunction.create(method, ""), beamCatalog.getUserFunctionDefinitions().javaScalarFunctions().get(ImmutableList.of("increment")));
}

Example 10 with JdbcConnection

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

the class BeamZetaSqlCatalogTest method rejectsScalarFunctionImplWithUnsupportedReturnType.

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