diff --git a/core/src/main/java/org/apache/calcite/adapter/jdbc/JdbcCatalogSchema.java b/core/src/main/java/org/apache/calcite/adapter/jdbc/JdbcCatalogSchema.java index e2e1478d13c9..5c80302d70c7 100644 --- a/core/src/main/java/org/apache/calcite/adapter/jdbc/JdbcCatalogSchema.java +++ b/core/src/main/java/org/apache/calcite/adapter/jdbc/JdbcCatalogSchema.java @@ -22,6 +22,7 @@ import org.apache.calcite.schema.Schema; import org.apache.calcite.schema.SchemaPlus; import org.apache.calcite.schema.Schemas; +import org.apache.calcite.schema.Wrapper; import org.apache.calcite.schema.impl.AbstractSchema; import org.apache.calcite.sql.SqlDialect; import org.apache.calcite.sql.SqlDialectFactory; @@ -54,7 +55,7 @@ * {@link JdbcSchema} for each schema name. Each JdbcSchema will populate its * tables on demand. */ -public class JdbcCatalogSchema extends AbstractSchema { +public class JdbcCatalogSchema extends AbstractSchema implements Wrapper { final DataSource dataSource; public final SqlDialect dialect; final JdbcConvention convention; @@ -137,6 +138,17 @@ public DataSource getDataSource() { return dataSource; } + + @Override public T unwrap(Class clazz) { + if (clazz.isInstance(this)) { + return clazz.cast(this); + } + if (clazz == DataSource.class) { + return clazz.cast(getDataSource()); + } + return null; + } + /** Contains sub-schemas by name, and the name of the default schema. */ private static class SubSchemaMap { final String defaultSchemaName; diff --git a/core/src/main/java/org/apache/calcite/adapter/jdbc/JdbcSchema.java b/core/src/main/java/org/apache/calcite/adapter/jdbc/JdbcSchema.java index 006390e3871b..438e99292c18 100644 --- a/core/src/main/java/org/apache/calcite/adapter/jdbc/JdbcSchema.java +++ b/core/src/main/java/org/apache/calcite/adapter/jdbc/JdbcSchema.java @@ -33,6 +33,7 @@ import org.apache.calcite.schema.SchemaVersion; import org.apache.calcite.schema.Schemas; import org.apache.calcite.schema.Table; +import org.apache.calcite.schema.Wrapper; import org.apache.calcite.sql.SqlDialect; import org.apache.calcite.sql.SqlDialectFactory; import org.apache.calcite.sql.SqlDialectFactoryImpl; @@ -75,7 +76,7 @@ * queries against this schema are executed against those tables, pushing down * as much as possible of the query logic to SQL. */ -public class JdbcSchema implements Schema { +public class JdbcSchema implements Schema, Wrapper { private static final Logger LOGGER = LoggerFactory.getLogger(JdbcSchema.class); final DataSource dataSource; @@ -512,6 +513,17 @@ protected Map getTypes() { return ImmutableSet.of(); } + @Override public T unwrap(Class clazz) { + if (clazz.isInstance(this)) { + return clazz.cast(this); + } + if (clazz == DataSource.class) { + return clazz.cast(getDataSource()); + } + return null; + } + + private static void close( @Nullable Connection connection, @Nullable Statement statement, diff --git a/core/src/main/java/org/apache/calcite/jdbc/CalciteSchema.java b/core/src/main/java/org/apache/calcite/jdbc/CalciteSchema.java index 4466bd3cd56b..76e1ba39f1c9 100644 --- a/core/src/main/java/org/apache/calcite/jdbc/CalciteSchema.java +++ b/core/src/main/java/org/apache/calcite/jdbc/CalciteSchema.java @@ -16,8 +16,6 @@ */ package org.apache.calcite.jdbc; -import org.apache.calcite.adapter.jdbc.JdbcCatalogSchema; -import org.apache.calcite.adapter.jdbc.JdbcSchema; import org.apache.calcite.linq4j.function.Experimental; import org.apache.calcite.linq4j.tree.Expression; import org.apache.calcite.materialize.Lattice; @@ -28,6 +26,7 @@ import org.apache.calcite.schema.SchemaVersion; import org.apache.calcite.schema.Table; import org.apache.calcite.schema.TableMacro; +import org.apache.calcite.schema.Wrapper; import org.apache.calcite.schema.impl.MaterializedViewTable; import org.apache.calcite.schema.impl.StarTable; import org.apache.calcite.util.NameMap; @@ -49,7 +48,6 @@ import java.util.NavigableMap; import java.util.NavigableSet; import java.util.Set; -import javax.sql.DataSource; import static java.util.Objects.requireNonNull; @@ -702,13 +700,8 @@ CalciteSchema calciteSchema() { if (clazz.isInstance(CalciteSchema.this.schema)) { return clazz.cast(CalciteSchema.this.schema); } - if (clazz == DataSource.class) { - if (schema instanceof JdbcSchema) { - return clazz.cast(((JdbcSchema) schema).getDataSource()); - } - if (schema instanceof JdbcCatalogSchema) { - return clazz.cast(((JdbcCatalogSchema) schema).getDataSource()); - } + if (schema instanceof Wrapper) { + return ((Wrapper) schema).unwrapOrThrow(clazz); } throw new ClassCastException("not a " + clazz); }