[CALCITE-6268] Support implementing custom JdbcSchema

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 extends Object> T unwrap(Class<T> 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;

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<String, RelProtoDataType> getTypes() {
     return ImmutableSet.of();
   }
 
+  @Override public <T extends Object> T unwrap(Class<T> 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,

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);
     }