From e0fb51ab5828e73255b80dc245d7c4b10fb15b11 Mon Sep 17 00:00:00 2001 From: Jan Blom Date: Tue, 30 Apr 2024 16:53:19 +0200 Subject: [PATCH] Switch to SQL for obtaining field metadata for Snowflake (default, JDBC can still be used through a system property or env.var) --- .../org/ohdsi/databases/DBConnection.java | 4 +++ .../org/ohdsi/databases/SnowflakeHandler.java | 14 +++++++- .../org/ohdsi/databases/StorageHandler.java | 35 ++++++++++++++----- .../ohdsi/whiterabbit/scan/ScanTestUtils.java | 4 +-- .../scan/SourceDataScanSnowflakeIT.java | 12 +++++++ 5 files changed, 57 insertions(+), 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 index 749c4e76..7dca1ecd 100644 --- 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 index 445bdc5a..e226940a 100644 --- 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 index 4192afca..ab85d5d1 100644 --- 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 getTableNames() { */ default List fetchTableStructure(String table, ScanParameters scanParameters) { List 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 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 index 1b286158..e1f07e41 100644 --- 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 index 06188ea9..9602cc39 100644 --- 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;