Switch to SQL for obtaining field metadata for Snowflake (default, JD…

…BC can still be used through a system property or env.var)
OHDSI · Apr 30, 2024 · e0fb51a · e0fb51a
1 parent e66b20f
commit e0fb51a
Show file tree

Hide file tree

Showing 5 changed files with 57 additions and 12 deletions.
diff --git a/rabbit-core/src/main/java/org/ohdsi/databases/DBConnection.java b/rabbit-core/src/main/java/org/ohdsi/databases/DBConnection.java
@@ -115,6 +115,10 @@ public void use(String database, DbType dbType) {
         }
     }
 
+    public QueryResult query(String sql) {
+        return new QueryResult(sql, this, verbose);
+    }
+
     public void execute(String sql) {
         execute(sql, false);
     }

diff --git a/rabbit-core/src/main/java/org/ohdsi/databases/SnowflakeHandler.java b/rabbit-core/src/main/java/org/ohdsi/databases/SnowflakeHandler.java
@@ -38,6 +38,7 @@
  */
 public enum SnowflakeHandler implements StorageHandler {
     INSTANCE();
+    public static final String WR_USE_SNOWFLAKE_JDBC_METADATA = "WR_USE_SNOWFLAKE_METADATA";
 
     DBConfiguration configuration = new SnowflakeConfiguration();
     private DBConnection snowflakeConnection = null;
@@ -117,7 +118,7 @@ private String resolveTableName(String tableName) {
     }
 
     @Override
-    public ResultSet getFieldNamesFromJDBC(String tableName) {
+    public ResultSet getFieldNamesInfo(String tableName) {
         try {
             String database = this.getDatabase();
             String schema = this.getSchema();
@@ -134,6 +135,17 @@ public ResultSet getFieldNamesFromJDBC(String tableName) {
         }
     }
 
+    @Override
+    public String getFieldsInformationQuery(String tableName) {
+        if (System.getenv(WR_USE_SNOWFLAKE_JDBC_METADATA) != null || System.getProperty(WR_USE_SNOWFLAKE_JDBC_METADATA) != null) {
+            return null;    // not providing a query forces use of JDBC metadata
+        } else {
+            return String.format(
+                    "SELECT column_name, data_type FROM %s.INFORMATION_SCHEMA.COLUMNS WHERE TABLE_SCHEMA = '%s' AND TABLE_NAME = '%s'",
+                    this.getDatabase().toUpperCase(), this.getSchema().toUpperCase(), tableName.toUpperCase());
+        }
+    }
+
     public String getTablesQuery(String database) {
         return String.format("SELECT TABLE_NAME FROM %s.INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = '%s'", this.getDatabase().toUpperCase(), this.getSchema().toUpperCase());
     }

diff --git a/rabbit-core/src/main/java/org/ohdsi/databases/StorageHandler.java b/rabbit-core/src/main/java/org/ohdsi/databases/StorageHandler.java
@@ -155,21 +155,38 @@ default List<String> getTableNames() {
      */
     default List<FieldInfo> fetchTableStructure(String table, ScanParameters scanParameters) {
         List<FieldInfo> fieldInfos = new ArrayList<>();
-        ResultSet rs = getFieldNamesFromJDBC(table);
-        try {
-            while (rs.next()) {
-                FieldInfo fieldInfo = new FieldInfo(scanParameters, rs.getString("COLUMN_NAME"));
-                fieldInfo.type = rs.getString("TYPE_NAME");
+        String fieldInfoQuery = getFieldsInformationQuery(table);
+        if (fieldInfoQuery != null) {
+            logger.warn("Obtaining field metadata through SQL");
+            QueryResult queryResult = getDBConnection().query(fieldInfoQuery);
+            for (Row row : queryResult) {
+                FieldInfo fieldInfo = new FieldInfo(scanParameters, row.getCells().get(0));
+                fieldInfo.type = row.getCells().get(1);
                 fieldInfo.rowCount = getTableSize(table);
                 fieldInfos.add(fieldInfo);
             }
-        } catch (
-                SQLException e) {
-            throw new RuntimeException(e.getMessage());
+        } else {
+            logger.warn("Obtaining field metadata through JDBC");
+            ResultSet rs = getFieldNamesInfo(table);
+            try {
+                while (rs.next()) {
+                    FieldInfo fieldInfo = new FieldInfo(scanParameters, rs.getString("COLUMN_NAME"));
+                    fieldInfo.type = rs.getString("TYPE_NAME");
+                    fieldInfo.rowCount = getTableSize(table);
+                    fieldInfos.add(fieldInfo);
+                }
+            } catch (
+                    SQLException e) {
+                throw new RuntimeException(e.getMessage());
+            }
         }
         return fieldInfos;
     }
 
+    default String getFieldsInformationQuery(String table) {
+        return null;
+    }
+
     /**
      * Retrieves column names (fields) for a table.
      *
@@ -179,7 +196,7 @@ default List<FieldInfo> fetchTableStructure(String table, ScanParameters scanPar
      * @param table name of the table to get the column names for
      * @return java.sql.ResultSet
      */
-    default ResultSet getFieldNamesFromJDBC(String table) {
+    default ResultSet getFieldNamesInfo(String table) {
         try {
             DatabaseMetaData metadata = getDBConnection().getMetaData();
             return metadata.getColumns(null, null, table, null);

diff --git a/whiterabbit/src/test/java/org/ohdsi/whiterabbit/scan/ScanTestUtils.java b/whiterabbit/src/test/java/org/ohdsi/whiterabbit/scan/ScanTestUtils.java
@@ -210,8 +210,8 @@ private static boolean matchTypeName(String type, String reference, DbType dbTyp
             case SNOWFLAKE:
                 switch (type) {
                     case "NUMBER": return reference.equals("integer") || reference.equals("numeric");
-                    case "VARCHAR": return reference.equals("character varying");
-                    case "TIMESTAMPNTZ": return reference.equals("timestamp without time zone");
+                    case "VARCHAR": case "TEXT": return reference.equals("character varying");
+                    case "TIMESTAMP_NTZ": case "TIMESTAMPNTZ": return reference.equals("timestamp without time zone");
                     default: throw new RuntimeException(String.format("Unsupported column type '%s' for DbType %s ", type, dbType.name()));
                 }
             case MYSQL:

diff --git a/whiterabbit/src/test/java/org/ohdsi/whiterabbit/scan/SourceDataScanSnowflakeIT.java b/whiterabbit/src/test/java/org/ohdsi/whiterabbit/scan/SourceDataScanSnowflakeIT.java
@@ -21,6 +21,7 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.io.TempDir;
+import org.ohdsi.databases.SnowflakeHandler;
 import org.ohdsi.databases.configuration.DbType;
 import org.ohdsi.whiterabbit.WhiteRabbitMain;
 import org.slf4j.Logger;
@@ -62,6 +63,17 @@ public void setUp() {
     }
 
     @Test
+    void testProcessSnowflakeFromIniWithSQLMetadata(@TempDir Path tempDir) throws URISyntaxException, IOException {
+        System.clearProperty(SnowflakeHandler.WR_USE_SNOWFLAKE_JDBC_METADATA);
+        testProcessSnowflakeFromIni(tempDir);
+    }
+
+    @Test
+    void testProcessSnowflakeFromIniWithJDBCMetadata(@TempDir Path tempDir) throws URISyntaxException, IOException {
+        System.setProperty(SnowflakeHandler.WR_USE_SNOWFLAKE_JDBC_METADATA, "true");
+        testProcessSnowflakeFromIni(tempDir);
+    }
+
     void testProcessSnowflakeFromIni(@TempDir Path tempDir) throws URISyntaxException, IOException {
         Assumptions.assumeTrue(new ScanTestUtils.PropertiesFileChecker("snowflake.env"), "Snowflake system properties file not available");
         Charset charset = StandardCharsets.UTF_8;