Revert "various cleanup fixes (#2467)"

This reverts commit 0ffb7b8.
awslabs · Dec 23, 2024 · ae5be3e · ae5be3e
1 parent dd1760c
commit ae5be3e
Show file tree

Hide file tree

Showing 6 changed files with 13 additions and 11 deletions.
diff --git a/...src/main/java/com/amazonaws/athena/connectors/db2as400/Db2As400EnvironmentProperties.java b/...src/main/java/com/amazonaws/athena/connectors/db2as400/Db2As400EnvironmentProperties.java
@@ -37,14 +37,14 @@ public Map<String, String> connectionPropertiesToEnvironment(Map<String, String>
         HashMap<String, String> environment = new HashMap<>();
 
         // now construct jdbc string
-        String connectionString = String.join("db2as400://jdbc:as400://", connectionProperties.get(HOST),
-                ";", connectionProperties.getOrDefault(JDBC_PARAMS, ""));
+        String connectionString = "db2as400://jdbc:as400://" + connectionProperties.get(HOST)
+                + ";" + connectionProperties.getOrDefault(JDBC_PARAMS, "");
 
         if (connectionProperties.containsKey(SECRET_NAME)) {
             if (connectionProperties.containsKey(JDBC_PARAMS)) { // need to add delimiter
                 connectionString = connectionString + ";";
             }
-            connectionString = String.join(connectionString, ":${", connectionProperties.get(SECRET_NAME), "}");
+            connectionString = connectionString + ":${" + connectionProperties.get(SECRET_NAME) + "}";
         }
 
         logger.debug("Constructed connection string: {}", connectionString);

diff --git a/athena-docdb/src/main/java/com/amazonaws/athena/connectors/docdb/DocDBMetadataHandler.java b/athena-docdb/src/main/java/com/amazonaws/athena/connectors/docdb/DocDBMetadataHandler.java
@@ -96,8 +96,6 @@ public class DocDBMetadataHandler
     //The Glue table property that indicates that a table matching the name of an DocDB table
     //is indeed enabled for use by this connector.
     private static final String DOCDB_METADATA_FLAG = "docdb-metadata-flag";
-    //The prefix of a connection string
-    protected static final String DOCDB_CONN_STRING_PREFIX = "mongodb://";
     //Used to filter out Glue tables which lack a docdb metadata flag.
     private static final TableFilter TABLE_FILTER = (Table table) -> table.parameters().containsKey(DOCDB_METADATA_FLAG);
     //The number of documents to scan when attempting to infer schema from an DocDB collection.
@@ -136,7 +134,7 @@ private MongoClient getOrCreateConn(MetadataRequest request)
     {
         String connStr = getConnStr(request);
         if (configOptions.containsKey(SECRET_NAME) && !hasEmbeddedSecret(connStr)) {
-            connStr = String.join(connStr.substring(0, DOCDB_CONN_STRING_PREFIX.length()), "${", configOptions.get(SECRET_NAME), "}@", connStr.substring(DOCDB_CONN_STRING_PREFIX.length()));
+            connStr = connStr.substring(0, 10) + "${" + configOptions.get(SECRET_NAME) + "}@" + connStr.substring(10);
         } 
         String endpoint = resolveSecrets(connStr);
         return connectionFactory.getOrCreateConn(endpoint);

diff --git a/athena-docdb/src/main/java/com/amazonaws/athena/connectors/docdb/DocDBRecordHandler.java b/athena-docdb/src/main/java/com/amazonaws/athena/connectors/docdb/DocDBRecordHandler.java
@@ -49,7 +49,6 @@
 import static com.amazonaws.athena.connector.lambda.handlers.GlueMetadataHandler.SOURCE_TABLE_PROPERTY;
 import static com.amazonaws.athena.connectors.docdb.DocDBFieldResolver.DEFAULT_FIELD_RESOLVER;
 import static com.amazonaws.athena.connectors.docdb.DocDBMetadataHandler.DOCDB_CONN_STR;
-import static com.amazonaws.athena.connectors.docdb.DocDBMetadataHandler.DOCDB_CONN_STRING_PREFIX;
 
 /**
  * Handles data read record requests for the Athena DocumentDB Connector.
@@ -111,7 +110,7 @@ private MongoClient getOrCreateConn(Split split)
             throw new RuntimeException(DOCDB_CONN_STR + " Split property is null! Unable to create connection.");
         }
         if (configOptions.containsKey(SECRET_NAME) && !hasEmbeddedSecret(connStr)) {
-            connStr = String.join(connStr.substring(0, DOCDB_CONN_STRING_PREFIX.length()), "${", configOptions.get(SECRET_NAME), "}@", connStr.substring(DOCDB_CONN_STRING_PREFIX.length()));
+            connStr = connStr.substring(0, 10) + "${" + configOptions.get(SECRET_NAME) + "}@" + connStr.substring(10);
         }
         String endpoint = resolveSecrets(connStr);
         return connectionFactory.getOrCreateConn(endpoint);

diff --git a/...na-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/JdbcEnvironmentProperties.java b/...na-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/JdbcEnvironmentProperties.java
@@ -39,8 +39,8 @@ public Map<String, String> connectionPropertiesToEnvironment(Map<String, String>
         HashMap<String, String> environment = new HashMap<>();
 
         // now construct jdbc string
-        String connectionString = String.join(getConnectionStringPrefix(connectionProperties), connectionProperties.get(HOST),
-                ":", connectionProperties.get(PORT), getDatabase(connectionProperties), getJdbcParameters(connectionProperties));
+        String connectionString = getConnectionStringPrefix(connectionProperties) + connectionProperties.get(HOST)
+                + ":" + connectionProperties.get(PORT) + getDatabase(connectionProperties) + getJdbcParameters(connectionProperties);
 
         environment.put(DEFAULT, connectionString);
         return environment;

diff --git a/...na-oracle/src/main/java/com/amazonaws/athena/connectors/oracle/OracleMetadataHandler.java b/...na-oracle/src/main/java/com/amazonaws/athena/connectors/oracle/OracleMetadataHandler.java
@@ -161,11 +161,13 @@ public void getPartitions(final BlockWriter blockWriter, final GetTableLayoutReq
         LOGGER.debug("{}: Schema {}, table {}", getTableLayoutRequest.getQueryId(), transformString(getTableLayoutRequest.getTableName().getSchemaName(), true),
                 transformString(getTableLayoutRequest.getTableName().getTableName(), true));
         try (Connection connection = getJdbcConnectionFactory().getConnection(getCredentialProvider())) {
-          List<String> parameters = Arrays.asList(transformString(getTableLayoutRequest.getTableName().getTableName(), true)); 
+          List<String> parameters = Arrays.asList(transformString(getTableLayoutRequest.getTableName().getTableName(), true));
+            //try (Statement statement = connection.createStatement(); ResultSet resultSet = statement.executeQuery(GET_PARTITIONS_QUERY + ))  
             try (PreparedStatement preparedStatement = new PreparedStatementBuilder().withConnection(connection).withQuery(GET_PARTITIONS_QUERY).withParameters(parameters).build();
                 ResultSet resultSet = preparedStatement.executeQuery()) {
                 // Return a single partition if no partitions defined
                 if (!resultSet.next()) {
+                    LOGGER.debug("here");
                     blockWriter.writeRows((Block block, int rowNum) -> {
                         LOGGER.debug("Parameters: " + BLOCK_PARTITION_COLUMN_NAME + " " + rowNum + " " + ALL_PARTITIONS);
                         block.setValue(BLOCK_PARTITION_COLUMN_NAME, rowNum, ALL_PARTITIONS);

diff --git a/...src/test/java/com/amazonaws/athena/connectors/snowflake/SnowflakeMetadataHandlerTest.java b/...src/test/java/com/amazonaws/athena/connectors/snowflake/SnowflakeMetadataHandlerTest.java
@@ -116,6 +116,7 @@ public void doGetTableLayout()
 
         PreparedStatement countsPreparedStatement = Mockito.mock(PreparedStatement.class);
         String GET_PKEY_COUNTS_QUERY = "SELECT \"pkey\", count(*) as COUNTS FROM \"testSchema\".\"testTable\" GROUP BY \"pkey\" ORDER BY COUNTS DESC";
+        System.err.println("test:" + GET_PKEY_COUNTS_QUERY);
         String[] countsColumns = new String[] {"pkey", SnowflakeMetadataHandler.COUNTS_COLUMN_NAME};
         Object[][] countsValues = {{"a", 1}};
         ResultSet countsResultSet = mockResultSet(countsColumns, countsValues, new AtomicInteger(-1));
@@ -180,6 +181,7 @@ public void doGetTableLayoutSinglePartition()
 
         PreparedStatement countsPreparedStatement = Mockito.mock(PreparedStatement.class);
         String GET_PKEY_COUNTS_QUERY = "SELECT \"pkey\", count(*) as COUNTS FROM \"testSchema\".\"testTable\" GROUP BY \"pkey\" ORDER BY COUNTS DESC";
+        System.err.println("test:" + GET_PKEY_COUNTS_QUERY);
         String[] countsColumns = new String[] {"pkey", SnowflakeMetadataHandler.COUNTS_COLUMN_NAME};
         Object[][] countsValues = {{"a", 1}};
         ResultSet countsResultSet = mockResultSet(countsColumns, countsValues, new AtomicInteger(-1));
@@ -223,6 +225,7 @@ public void doGetTableLayoutMaxPartition()
         long pageCount = (long) (Math.ceil(totalActualRecordCount / MAX_PARTITION_COUNT));
         long partitionActualRecordCount = (totalActualRecordCount <= 10000) ? (long) totalActualRecordCount : pageCount;
         double limit = (int) Math.ceil(totalActualRecordCount / partitionActualRecordCount);
+//        double limit = 1;
         long offset = 0;
         String[] columns = {"partition"};
         int[] types = {Types.VARCHAR};