Add passing through metadataMap to CalciteMetaTable constructor

looker-open-source · Aug 22, 2023 · cefef06 · cefef06
1 parent 76e7663
commit cefef06
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 3 deletions.
diff --git a/core/src/main/java/org/apache/calcite/jdbc/CalciteMetaImpl.java b/core/src/main/java/org/apache/calcite/jdbc/CalciteMetaImpl.java
@@ -75,6 +75,7 @@
 import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -277,7 +278,10 @@ private static ImmutableMap.Builder<DatabaseProperty, Object> addProperty(
         "TYPE_SCHEM",
         "TYPE_NAME",
         "SELF_REFERENCING_COL_NAME",
-        "REF_GENERATION");
+        "REF_GENERATION",
+        "EXPLORE_LABEL",
+        "EXPLORE_DESCRIPTION",
+        "EXPLORE_TAGS");
   }
 
   @Override public MetaResultSet getTypeInfo(ConnectionHandle ch) {
@@ -386,21 +390,27 @@ Enumerable<MetaTable> tables(final MetaSchema schema_) {
               requireNonNull(schema.calciteSchema.getTable(name, true),
                   () -> "table " + name + " is not found (case sensitive)")
                   .getTable();
+          HashMap<String, String> metadataMap = (table.getTableMetadata() != null)
+              ? table.getTableMetadata() : new HashMap<>();
           return new CalciteMetaTable(table,
               schema.tableCatalog,
               schema.tableSchem,
-              name);
+              name,
+              metadataMap);
         })
         .concat(
             Linq4j.asEnumerable(
                 schema.calciteSchema.getTablesBasedOnNullaryFunctions()
                     .entrySet())
                 .select(pair -> {
                   final Table table = pair.getValue();
+                  HashMap<String, String> metadataMap = (table.getTableMetadata() != null)
+                      ? table.getTableMetadata() : new HashMap<>();
                   return new CalciteMetaTable(table,
                       schema.tableCatalog,
                       schema.tableSchem,
-                      pair.getKey());
+                      pair.getKey(),
+                      metadataMap);
                 }));
   }
 
@@ -831,6 +841,12 @@ private static class CalciteMetaTable extends MetaTable {
           calciteTable.getJdbcTableType().jdbcName);
       this.calciteTable = requireNonNull(calciteTable, "calciteTable");
     }
+    CalciteMetaTable(Table calciteTable, String tableCat,
+        String tableSchem, String tableName, HashMap<String, String> metadataMap) {
+      super(tableCat, tableSchem, tableName,
+          calciteTable.getJdbcTableType().jdbcName, metadataMap);
+      this.calciteTable = requireNonNull(calciteTable, "calciteTable");
+    }
   }
 
   /** Metadata describing a Calcite schema. */

diff --git a/core/src/main/java/org/apache/calcite/schema/Table.java b/core/src/main/java/org/apache/calcite/schema/Table.java
@@ -24,6 +24,8 @@
 
 import org.checkerframework.checker.nullness.qual.Nullable;
 
+import java.util.HashMap;
+
 /**
  * Table.
  *
@@ -82,4 +84,15 @@ public interface Table {
    */
   boolean rolledUpColumnValidInsideAgg(String column, SqlCall call,
       @Nullable SqlNode parent, @Nullable CalciteConnectionConfig config);
+
+  /**
+   * Gets Looker Explore metadata if available.
+   */
+  default HashMap<String, String> getTableMetadata() {
+    HashMap<String, String> map = new HashMap<>();
+    map.put("EXPLORE_LABEL", null);
+    map.put("EXPLORE_DESCRIPTION", null);
+    map.put("EXPLORE_TAGS", null);
+    return map;
+  }
 }
diff --git a/core/src/test/java/org/apache/calcite/jdbc/CalciteRemoteDriverTest.java b/core/src/test/java/org/apache/calcite/jdbc/CalciteRemoteDriverTest.java
@@ -129,6 +129,17 @@ class CalciteRemoteDriverTest {
         }
       };
 
+  private static final Function<Connection, ResultSet> GET_TABLES =
+      connection -> {
+        try {
+          return connection.getMetaData().getTables(connection.getCatalog(), null, null,
+              new String[] {"TABLE"});
+        } catch (SQLException e) {
+          throw TestUtil.rethrow(e);
+        }
+      };
+
+
   private static final Function<Connection, ResultSet> GET_TYPEINFO =
       connection -> {
         try {
@@ -284,6 +295,22 @@ protected static Connection getRemoteConnection() throws SQLException {
         .returns(CalciteAssert.checkResultContains("COLUMN_NAME=EMPNO"));
   }
 
+  @Test void testRemoteTables() throws SQLException {
+    final Connection connection =
+        DriverManager.getConnection("jdbc:avatica:remote:factory=" + LJS);
+    assertThat(connection.isClosed(), is(false));
+    final ResultSet resultSet =
+        connection.getMetaData().getTables(connection.getCatalog(), null, null, new String[] {"TABLE"});
+    final ResultSetMetaData metaData = resultSet.getMetaData();
+    assertThat(metaData.getColumnCount(), is(13));
+    assertThat(metaData.getColumnName(11), is("EXPLORE_LABEL"));
+    assertThat(metaData.getColumnName(12), is("EXPLORE_DESCRIPTION"));
+    assertThat(metaData.getColumnName(13), is("EXPLORE_TAGS"));
+    resultSet.close();
+    connection.close();
+    assertThat(connection.isClosed(), is(true));
+  }
+
   @Test void testRemoteTypeInfo() {
     // TypeInfo does not include internal types (NULL, SYMBOL, ANY, etc.)
     CalciteAssert.hr()