Fixing other apis in snowflake connector (Netflix#256)

Author: zhljen

metacat-connector-jdbc/src/main/java/com/netflix/metacat/connector/jdbc/services/JdbcConnectorTableService.java

@@ -65,9 +65,9 @@ public class JdbcConnectorTableService implements ConnectorTableService {
     private static final char LEFT_PAREN = '(';
     private static final char RIGHT_PAREN = ')';
     private static final char SPACE = ' ';
-    private final DataSource dataSource;
-    private final JdbcTypeConverter typeConverter;
+    protected final DataSource dataSource;
     private final JdbcExceptionMapper exceptionMapper;
+    private final JdbcTypeConverter typeConverter;
 
     /**
      * Constructor.
@@ -95,8 +95,7 @@ public void delete(@Nonnull final ConnectorRequestContext context, @Nonnull fina
         final String databaseName = name.getDatabaseName();
         final String tableName = name.getTableName();
         log.debug("Attempting to delete table {} from database {} for request {}", tableName, databaseName, context);
-        try (Connection connection = this.dataSource.getConnection()) {
-            connection.setSchema(databaseName);
+        try (Connection connection = this.getConnection(name.getDatabaseName())) {
             JdbcConnectorUtils.executeUpdate(connection, this.getDropTableSql(name, tableName));
             log.debug("Deleted table {} from database {} for request {}", tableName, databaseName, context);
         } catch (final SQLException se) {
@@ -111,9 +110,7 @@ public void delete(@Nonnull final ConnectorRequestContext context, @Nonnull fina
     public TableInfo get(@Nonnull final ConnectorRequestContext context, @Nonnull final QualifiedName name) {
         log.debug("Beginning to get table metadata for qualified name {} for request {}", name, context);
 
-        try (Connection connection = this.dataSource.getConnection()) {
-            final String database = name.getDatabaseName();
-            connection.setSchema(database);
+        try (Connection connection = this.getConnection(name.getDatabaseName())) {
             final ImmutableList.Builder<FieldInfo> fields = ImmutableList.builder();
             try (ResultSet columns = this.getColumns(connection, name)) {
                 while (columns.next()) {
@@ -196,8 +193,7 @@ public List<QualifiedName> listNames(
         final String catalog = name.getCatalogName();
         final String database = name.getDatabaseName();
 
-        try (Connection connection = this.dataSource.getConnection()) {
-            connection.setSchema(database);
+        try (Connection connection = this.getConnection(database)) {
             final List<QualifiedName> names = Lists.newArrayList();
             try (ResultSet tables = this.getTables(connection, name, prefix, null)) {
                 while (tables.next()) {
@@ -248,7 +244,7 @@ public void rename(
                 "Database names must match and they are " + oldDatabaseName + " and " + newDatabaseName
             );
         }
-        try (Connection connection = this.dataSource.getConnection()) {
+        try (Connection connection = this.getConnection(oldDatabaseName)) {
             connection.setSchema(oldDatabaseName);
             JdbcConnectorUtils.executeUpdate(
                 connection,
@@ -267,6 +263,12 @@ public void rename(
         }
     }
 
+    protected Connection getConnection(@Nonnull @NonNull final String schema) throws SQLException {
+        final Connection connection = this.dataSource.getConnection();
+        connection.setSchema(schema);
+        return connection;
+    }
+
     @Override
     public boolean exists(@Nonnull final ConnectorRequestContext context, @Nonnull final QualifiedName name) {
         boolean result = false;
@@ -398,7 +400,6 @@ protected String buildSourceType(
 
     /**
      * Build the SQL for renaming a table out of the components provided. SQL will be executed.
-     *
      * @param oldName           The fully qualified name for the current table
      * @param finalOldTableName The string for what the current table should be called in the sql
      * @param finalNewTableName The string for what the new name fo the table should be in the sql

metacat-connector-snowflake/src/main/java/com/netflix/metacat/connector/snowflake/SnowflakeConnectorTableService.java

@@ -55,6 +55,7 @@ public class SnowflakeConnectorTableService extends JdbcConnectorTableService {
     private static final String SQL_GET_AUDIT_INFO
         = "select created, last_altered from information_schema.tables"
         + " where table_catalog=? and table_schema=? and table_name=?";
+
     /**
      * Constructor.
      *
@@ -130,6 +131,13 @@ public void rename(@Nonnull final ConnectorRequestContext context,
                        @Nonnull final QualifiedName newName) {
         super.rename(context, getSnowflakeName(oldName), getSnowflakeName(newName));
     }
+    
+    @Override
+    protected Connection getConnection(@Nonnull @NonNull final String schema) throws SQLException {
+        final Connection connection = this.dataSource.getConnection();
+        connection.setSchema(connection.getCatalog());
+        return connection;
+    }
 
     /**
      * {@inheritDoc}
@@ -139,6 +147,19 @@ public boolean exists(@Nonnull final ConnectorRequestContext context, @Nonnull f
         return super.exists(context, getSnowflakeName(name));
     }
 
+    @Override
+    protected ResultSet getColumns(
+        @Nonnull @NonNull final Connection connection,
+        @Nonnull @NonNull final QualifiedName name
+    ) throws SQLException {
+        return connection.getMetaData().getColumns(
+            connection.getCatalog(),
+            name.getDatabaseName(),
+            name.getTableName(),
+            JdbcConnectorUtils.MULTI_CHARACTER_SEARCH
+        );
+    }
+
     /**
      * {@inheritDoc}
      */

metacat-connector-snowflake/src/test/groovy/com/netflix/metacat/connector/snowflake/SnowflakeConnectorTableServiceSpec.groovy

@@ -63,7 +63,8 @@ class SnowflakeConnectorTableServiceSpec extends Specification {
         then:
         1 * this.dataSource.getConnection() >> connection
         1 * connection.createStatement() >> statement
-        1 * connection.setSchema(database)
+        1 * connection.getCatalog() >> "dse"
+        1 * connection.setSchema(_)
         1 * statement.executeUpdate(
             "DROP TABLE " + table
         )
@@ -83,7 +84,8 @@ class SnowflakeConnectorTableServiceSpec extends Specification {
         then:
         1 * this.dataSource.getConnection() >> connection
         1 * connection.createStatement() >> statement
-        1 * connection.setSchema(database)
+        1 * connection.getCatalog() >> "dse"
+        1 * connection.setSchema(_)
         1 * statement.executeUpdate(
             "DROP TABLE " + table
         )
@@ -112,6 +114,8 @@ class SnowflakeConnectorTableServiceSpec extends Specification {
         then:
         1 * this.dataSource.getConnection() >> connection
         1 * connection.createStatement() >> statement
+        1 * connection.getCatalog() >> "dse"
+        2 * connection.setSchema(_)
         1 * statement.executeUpdate(
             "ALTER TABLE C RENAME TO D"
         )
@@ -129,6 +133,8 @@ class SnowflakeConnectorTableServiceSpec extends Specification {
         then:
         1 * this.dataSource.getConnection() >> connection
         1 * connection.createStatement() >> statement
+        1 * connection.getCatalog() >> "dse"
+        2 * connection.setSchema(_)
         1 * statement.executeUpdate(
             "ALTER TABLE C RENAME TO D"
         )