Refactoring, addressing Michaels comments

Author: akuzin1

athena-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/manager/JdbcRecordHandler.java

@@ -82,7 +82,6 @@
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
-import java.sql.ResultSetMetaData;
 import java.sql.SQLDataException;
 import java.sql.SQLException;
 import java.sql.Struct;
@@ -142,21 +141,14 @@ protected JdbcCredentialProvider getCredentialProvider()
         return null;
     }
 
-    protected java.util.Optional<Boolean> getAutoCommit()
-    {
-      return java.util.Optional.of(false);
-    }
-
     @Override
     public void readWithConstraint(BlockSpiller blockSpiller, ReadRecordsRequest readRecordsRequest, QueryStatusChecker queryStatusChecker)
             throws Exception
     {
         LOGGER.info("Read Record Request {}: Catalog: {}, table {}, splits {}", readRecordsRequest.getQueryId(), readRecordsRequest.getCatalogName(), readRecordsRequest.getTableName(),
                 readRecordsRequest.getSplit().getProperties());
         try (Connection connection = this.jdbcConnectionFactory.getConnection(getCredentialProvider())) {
-            if (this.getAutoCommit().isPresent()) {
-                connection.setAutoCommit(this.getAutoCommit().get()); // For consistency. This is needed to be false to enable streaming for some database types.
-            }
+            connection.setAutoCommit(false); // For consistency. This is needed to be false to enable streaming for some database types.
             try (PreparedStatement preparedStatement = buildSplitSql(connection, readRecordsRequest.getCatalogName(), readRecordsRequest.getTableName(),
                     readRecordsRequest.getSchema(), readRecordsRequest.getConstraints(), readRecordsRequest.getSplit());
                     ResultSet resultSet = preparedStatement.executeQuery()) {
@@ -185,9 +177,7 @@ else if (next.getType() instanceof ArrowType.Struct) {
                     rowsReturnedFromDatabase++;
                 }
                 LOGGER.info("{} rows returned by database.", rowsReturnedFromDatabase);
-                if (this.getAutoCommit().isPresent()) {
-                    connection.commit();
-                }
+                connection.commit();
             }
         }
     }
@@ -211,29 +201,6 @@ protected FieldWriterFactory makeFactory(Field field)
                 };
     }
 
-    public static void printResultSet(ResultSet rs) throws SQLException
-    {
-        LOGGER.info("Printing Record Handler Result Set:");
-        ResultSetMetaData rsmd = rs.getMetaData();
-        int columns = rsmd.getColumnCount();
-    
-        // Print the column names
-        String columnNames = "";
-        for (int i = 1; i <= columns; i++) {
-            columnNames += rsmd.getColumnName(i) + ", ";
-        }
-        LOGGER.info(columnNames);
-    
-        // Print the column values
-        while (rs.next()) {
-            String newColumn = "";
-            for (int i = 1; i <= columns; i++) {
-                newColumn += rs.getObject(i) + ", ";
-            }
-            LOGGER.info(newColumn);
-        }
-    }
-
     /**
      * Create a field extractor for complex Nested type.
      * @param field Field's metadata information.

athena-opensearch/athena-opensearch.yaml

@@ -17,6 +17,9 @@ Parameters:
     Description: 'This is the name of the lambda function that will be created. This name must satisfy the pattern ^[a-z0-9-_]{1,64}$'
     Type: String
     AllowedPattern: ^[a-z0-9-_]{1,64}$
+  GlueDatabaseName:
+    Description: 'Name of Glue Database that contains schemas of indices in OpenSearch Domain, also used as Database name.'
+    Type: String
   DefaultConnectionString:
     Description: 'The default connection string is used when catalog is "lambda:${LambdaFunctionName}". Catalog specific Connection Strings can be added later. Format: ${DatabaseType}://${NativeJdbcConnectionString}.'
     Type: String
@@ -68,6 +71,7 @@ Resources:
           disable_spill_encryption: !Ref DisableSpillEncryption
           spill_bucket: !Ref SpillBucket
           spill_prefix: !Ref SpillPrefix
+          glue_database: !Ref GlueDatabaseName
           default: !Ref DefaultConnectionString
       FunctionName: !Ref LambdaFunctionName
       Handler: "com.amazonaws.athena.connectors.opensearch.OpensearchMuxCompositeHandler"

athena-opensearch/src/main/java/com/amazonaws/athena/connectors/opensearch/OpensearchConstants.java

@@ -25,6 +25,7 @@ public final class OpensearchConstants
     public static final String OPENSEARCH_DRIVER_CLASS = "org.opensearch.jdbc.Driver";
     public static final int OPENSEARCH_DEFAULT_PORT = 443;
     public static final String OPENSEARCH_QUOTE_CHARACTER = "`";
+    public static final int OPENSEARCH_FETCH_SIZE = 10000;
 
     private OpensearchConstants() {}
 }

athena-opensearch/src/main/java/com/amazonaws/athena/connectors/opensearch/OpensearchGlueHandler.java

@@ -68,11 +68,6 @@ public void getPartitions(final BlockWriter blockWriter, final GetTableLayoutReq
 
     public boolean isDisabled()
     {
-        if (super.getAwsGlue() == null) {
-            return true;
-        }
-        else {
-            return false;
-        }
+        return super.getAwsGlue() == null;
     }
 }

athena-opensearch/src/main/java/com/amazonaws/athena/connectors/opensearch/OpensearchMetadataHandler.java

@@ -107,10 +107,11 @@ public class OpensearchMetadataHandler
     private static final Logger LOGGER = LoggerFactory.getLogger(OpensearchMetadataHandler.class);
     private static final int MAX_SPLITS_PER_REQUEST = 1000_000;
     private final String schemaName;
-    private static final String DATA_STREAM_REGEX = "(\\.ds-)(.+)(-\\d\\d\\d\\d\\d\\d)";
+    // data streams follow the following naming convention: .ds-<data-stream>-<generation> where generation is a six-digit, zero-padded integer
+    private static final String DATA_STREAM_REGEX = "(\\.ds-)(.+)(-\\d{6})";
 
     private OpensearchGlueHandler glueHandler;
-    private final boolean awsGlue;
+    private final boolean isGlueDisabled;
 
     /**
      * Instantiates handler to be used by Lambda function directly.
@@ -133,10 +134,10 @@ public OpensearchMetadataHandler(DatabaseConnectionConfig databaseConnectionConf
     public OpensearchMetadataHandler(DatabaseConnectionConfig databaseConnectionConfig, JdbcConnectionFactory jdbcConnectionFactory, java.util.Map<String, String> configOptions)
     {
         super(databaseConnectionConfig, jdbcConnectionFactory, configOptions);
-        this.schemaName = configOptions.getOrDefault("schema_name", "default");
+        this.schemaName = configOptions.getOrDefault("glue_database", "default");
 
         this.glueHandler = new OpensearchGlueHandler(configOptions);
-        this.awsGlue = glueHandler.isDisabled();
+        this.isGlueDisabled = glueHandler.isDisabled();
     }
 
     @VisibleForTesting
@@ -148,10 +149,10 @@ protected OpensearchMetadataHandler(
         java.util.Map<String, String> configOptions)
     {
         super(databaseConnectionConfig, secretsManager, athena, jdbcConnectionFactory, configOptions);
-        this.schemaName = configOptions.getOrDefault("schema_name", "default");
+        this.schemaName = configOptions.getOrDefault("glue_database", "default");
 
         this.glueHandler = new OpensearchGlueHandler(configOptions);
-        this.awsGlue = glueHandler.isDisabled();
+        this.isGlueDisabled = glueHandler.isDisabled();
     }
 
     @Override
@@ -259,7 +260,7 @@ public GetTableResponse doGetTable(final BlockAllocator blockAllocator, final Ge
 
         // Look at GLUE catalog first.
         try {
-            if (!this.awsGlue) {
+            if (!this.isGlueDisabled) {
                 schema = glueHandler.doGetTable(blockAllocator, getTableRequest).getSchema();
                 LOGGER.info("glueDoGetTable: Retrieved schema for table[{}] from AWS Glue.", getTableRequest.getTableName());
             }
@@ -369,22 +370,16 @@ private List<TableName> listTables(final Connection jdbcConnection, final String
             throws SQLException
     {
         try (ResultSet resultSet = getTables(jdbcConnection, databaseName)) {
-            ImmutableList.Builder<TableName> list = ImmutableList.builder();
+            HashSet<TableName> list = new HashSet<TableName>();
             while (resultSet.next()) {
                 // checks if table name is internal table ex: .kibana
                 String tableName = resultSet.getString("TABLE_NAME");
-                if (tableName.startsWith(".")) {
-                    if (!tableName.matches(DATA_STREAM_REGEX)) {
-                        continue;
-                    }
+                if (tableName.startsWith(".") && !tableName.matches(DATA_STREAM_REGEX)) {
+                    continue;
                 }
-                // update schema from null because Opensearch doesn't have schema but change later
-                // resultSet.updateString("TABLE_SCHEM", databaseName);
-                
-                // Temporary in schemaName 
                 list.add(getSchemaTableName(resultSet, databaseName));
             }
-            return list.build();
+            return new ImmutableList.builder().addAll(list).build();
         }
     }
 
@@ -393,6 +388,9 @@ private TableName getSchemaTableName(final ResultSet resultSet, String schemaNam
     {
         String tableName = resultSet.getString("TABLE_NAME");
         String tableSchema = resultSet.getString("TABLE_SCHEM");
+        // Data streams follow the following naming convention: (.ds-)(<data-stream>)(-<generation>) where generation is a six-digit, zero-padded integer 
+        // and each paranthesis signifies matching groups.
+        // Therefore I use the second matching group as table name because it represents the single resource name that stores time series data across multiple indices
         Pattern pattern = Pattern.compile(DATA_STREAM_REGEX);
         Matcher matcher = pattern.matcher(tableName);
         if (matcher.find()) {

athena-opensearch/src/main/java/com/amazonaws/athena/connectors/opensearch/OpensearchUtils.java athena-opensearch/src/main/java/com/amazonaws/athena/connectors/opensearch/OpensearchNestedStructBuilder.java

@@ -28,9 +28,9 @@
 import java.util.Arrays;
 import java.util.List;
 
-public class OpensearchUtils
+public class OpensearchNestedStructBuilder
 {
-    private OpensearchUtils()
+    private OpensearchNestedStructBuilder()
     {
     }
 

athena-opensearch/src/main/java/com/amazonaws/athena/connectors/opensearch/OpensearchRecordHandler.java

@@ -48,7 +48,8 @@
 import static com.amazonaws.athena.connectors.opensearch.OpensearchConstants.OPENSEARCH_DEFAULT_PORT;
 import static com.amazonaws.athena.connectors.opensearch.OpensearchConstants.OPENSEARCH_DRIVER_CLASS;
 import static com.amazonaws.athena.connectors.opensearch.OpensearchConstants.OPENSEARCH_NAME;
-import static com.amazonaws.athena.connectors.opensearch.OpensearchConstants.OPENSEARCH_QUOTE_CHARACTER;;
+import static com.amazonaws.athena.connectors.opensearch.OpensearchConstants.OPENSEARCH_QUOTE_CHARACTER;
+import static com.amazonaws.athena.connectors.opensearch.OpensearchConstants.OPENSEARCH_FETCH_SIZE;
 
 
 /**
@@ -59,8 +60,6 @@ public class OpensearchRecordHandler
 {
     private static final Logger LOGGER = LoggerFactory.getLogger(OpensearchRecordHandler.class);
 
-    private static final String OPENSEARCH_QUOTE_CHARACTER = "`";
-
     private final JdbcSplitQueryBuilder jdbcSplitQueryBuilder;
 
     /**
@@ -98,14 +97,8 @@ public PreparedStatement buildSplitSql(Connection jdbcConnection, String catalog
     {
         PreparedStatement preparedStatement = jdbcSplitQueryBuilder.buildSql(jdbcConnection, null, null, tableName.getTableName(), schema, constraints, split);
 
-        // Disable fetching all rows.
-        preparedStatement.setFetchSize(10000);
+        // by setting fetch size we explicitly force pagination to happen, otherwise it will fall back to legacy engine and only return 200 rows
+        preparedStatement.setFetchSize(OPENSEARCH_FETCH_SIZE);
         return preparedStatement;
     }
-
-    @Override
-    protected java.util.Optional<Boolean> getAutoCommit()
-    {
-      return java.util.Optional.empty();
-    }
 }