added new connector for OpenSearch data source

Author: akuzin1

athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/data/BlockUtils.java

@@ -849,6 +849,12 @@ else if (value instanceof Integer) {
                     if (value == null) {
                         float4Writer.writeNull();
                     }
+                    else if (value instanceof java.lang.Integer) {
+                        float4Writer.writeFloat4((int) value);
+                    }
+                    else if (value instanceof java.lang.Double) {
+                        float4Writer.writeFloat4(((Double) value).floatValue());
+                    }
                     else {
                         float4Writer.writeFloat4((float) value);
                     }
@@ -926,6 +932,9 @@ else if (value != null && value instanceof Long) {
                     if (value == null) {
                         bigIntWriter.writeNull();
                     }
+                    else if (value instanceof java.lang.Integer) {
+                        bigIntWriter.writeBigInt(((Integer) value).longValue());
+                    }
                     else {
                         bigIntWriter.writeBigInt((long) value);
                     }

athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/data/FieldBuilder.java

@@ -42,6 +42,7 @@ public class FieldBuilder
     private final boolean nullable;
     //Using LinkedHashMap because Apache Arrow makes field order important so honoring that contract here
     private final Map<String, Field> children = new LinkedHashMap<>();
+    private final Map<String, FieldBuilder> nestedChildren = new LinkedHashMap<>();
 
     /**
      * Creates a FieldBuilder for a Field with the given name and type.
@@ -96,6 +97,12 @@ public FieldBuilder addField(String fieldName, ArrowType type, List<Field> child
         return this;
     }
 
+    public FieldBuilder addFieldBuilder(String fieldName, FieldBuilder child)
+    {
+        this.nestedChildren.put(fieldName, child);
+        return this;
+    }
+
     public FieldBuilder addField(String fieldName, ArrowType type, boolean nullable, List<Field> children)
     {
         this.children.put(fieldName, new Field(fieldName, new FieldType(nullable, type, null), children));
@@ -263,18 +270,31 @@ public FieldBuilder addDateMilliField(String fieldName)
         return this;
     }
 
+    public String getName()
+    {
+        return this.name;
+    }
+
     public Field getChild(String fieldName)
     {
         return children.get(fieldName);
     }
 
+    public FieldBuilder getNestedChild(String fieldName)
+    {
+        return nestedChildren.get(fieldName);
+    }
+
     /**
      * Builds the fields.
      *
      * @return The newly constructed Field.
      */
     public Field build()
     {
+        for (Map.Entry<String, FieldBuilder> next : nestedChildren.entrySet()) {
+            children.put(next.getKey(), next.getValue().build());
+        }
         return new Field(name, new FieldType(nullable, type, null), new ArrayList<>(children.values()));
     }
 }

athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/data/SchemaBuilder.java

@@ -89,6 +89,18 @@ public SchemaBuilder addStructField(String fieldName)
         return this;
     }
 
+    /**
+     * Adds a new Nested STRUCT Field to the Schema as a top-level Field.
+     *
+     * @param fieldName The name of the field to add.
+     * @return This SchemaBuilder itself.
+     */
+    public SchemaBuilder addNestedField(String fieldName, FieldBuilder newField)
+    {
+        nestedFieldBuilderMap.put(fieldName, newField);
+        return this;
+    }
+
     /**
      * Adds a new LIST Field to the Schema as a top-level Field.
      *

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

@@ -76,7 +76,7 @@ public abstract class JdbcMetadataHandler
     private static final Logger LOGGER = LoggerFactory.getLogger(JdbcMetadataHandler.class);
     private static final String SQL_SPLITS_STRING = "select min(%s), max(%s) from %s.%s;";
     private static final int DEFAULT_NUM_SPLITS = 20;
-    private final JdbcConnectionFactory jdbcConnectionFactory;
+    protected final JdbcConnectionFactory jdbcConnectionFactory;
     private final DatabaseConnectionConfig databaseConnectionConfig;
     private final SplitterFactory splitterFactory = new SplitterFactory();
 
@@ -275,7 +275,7 @@ private ResultSet getColumns(final String catalogName, final TableName tableHand
         return metadata.getColumns(
                 catalogName,
                 escapeNamePattern(tableHandle.getSchemaName(), escape),
-                escapeNamePattern(tableHandle.getTableName(), escape),
+                escapeNamePattern(tableHandle.getTableName(), null), 
                 null);
     }
 

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

@@ -76,13 +76,20 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import javax.swing.RowFilter.Entry;
+
 import java.sql.Array;
 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;
+import java.sql.Timestamp;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
@@ -135,14 +142,21 @@ 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("{}: Catalog: {}, table {}, splits {}", readRecordsRequest.getQueryId(), readRecordsRequest.getCatalogName(), readRecordsRequest.getTableName(),
+        LOGGER.info("Read Record Request {}: Catalog: {}, table {}, splits {}", readRecordsRequest.getQueryId(), readRecordsRequest.getCatalogName(), readRecordsRequest.getTableName(),
                 readRecordsRequest.getSplit().getProperties());
         try (Connection connection = this.jdbcConnectionFactory.getConnection(getCredentialProvider())) {
-            connection.setAutoCommit(false); // For consistency. This is needed to be false to enable streaming for some database types.
+            if (this.getAutoCommit().isPresent()) {
+                connection.setAutoCommit(this.getAutoCommit().get()); // 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()) {
@@ -153,6 +167,9 @@ public void readWithConstraint(BlockSpiller blockSpiller, ReadRecordsRequest rea
                     if (next.getType() instanceof ArrowType.List) {
                         rowWriterBuilder.withFieldWriterFactory(next.getName(), makeFactory(next));
                     }
+                    else if (next.getType() instanceof ArrowType.Struct) {
+                        rowWriterBuilder.withFieldWriterFactory(next.getName(), makeStructFactory(next));
+                    }
                     else {
                         rowWriterBuilder.withExtractor(next.getName(), makeExtractor(next, resultSet, partitionValues));
                     }
@@ -168,8 +185,9 @@ public void readWithConstraint(BlockSpiller blockSpiller, ReadRecordsRequest rea
                     rowsReturnedFromDatabase++;
                 }
                 LOGGER.info("{} rows returned by database.", rowsReturnedFromDatabase);
-
-                connection.commit();
+                if (this.getAutoCommit().isPresent()) {
+                    connection.commit();
+                }
             }
         }
     }
@@ -193,6 +211,54 @@ 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.
+     * @return Extractor for the List type.
+     */
+    protected FieldWriterFactory makeStructFactory(Field field)
+    {
+        return (FieldVector vector, Extractor extractor, ConstraintProjector constraint) ->
+                (FieldWriter) (Object context, int rowNum) ->
+                {
+                    Object nestedObj = ((ResultSet) context).getObject(field.getName());
+                    if (!((ResultSet) context).wasNull() && nestedObj instanceof Struct) {
+                        Map<String, Object> renestedMap = new HashMap<>();
+                        for (Object obj : ((Struct) nestedObj).getAttributes()) {
+                            Map.Entry<?, ?> entry = (Map.Entry<?, ?>) obj;
+                            String key = entry.getKey().toString();
+                            Object value = entry.getValue();
+                            renestedMap.put(key, value);
+                        }
+                        BlockUtils.setComplexValue(vector, rowNum, FieldResolver.DEFAULT, renestedMap);
+                    }
+                    return true;
+                };
+    }
+
     /**
      * Creates an Extractor for the given field. In this example the extractor just creates some random data.
      */
@@ -278,8 +344,17 @@ protected Extractor makeExtractor(Field field, ResultSet resultSet, Map<String,
             case DATEMILLI:
                 return (DateMilliExtractor) (Object context, NullableDateMilliHolder dst) ->
                 {
-                    if (resultSet.getTimestamp(fieldName) != null) {
-                        dst.value = resultSet.getTimestamp(fieldName).getTime();
+                    //try catch needed for OpenSearch type date which actually is returned as timestamp
+                    try {
+                        if (resultSet.getTimestamp(fieldName) != null) {
+                            dst.value = resultSet.getTimestamp(fieldName).getTime();
+                        }
+                    }
+                    catch (SQLDataException e) {
+                        //OpenSearch returns it as type String
+                        if (resultSet.getString(fieldName) != null) {
+                            dst.value = Timestamp.valueOf(resultSet.getString(fieldName)).getTime();
+                        }
                     }
                     dst.isSet = resultSet.wasNull() ? 0 : 1;
                 };

athena-jdbc/src/test/java/com/amazonaws/athena/connectors/jdbc/TestBase.java

@@ -24,6 +24,7 @@
 
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.Struct;
 import java.util.Arrays;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -78,6 +79,26 @@ else if (argument instanceof String) {
         });
 
 
+        Mockito.when(resultSet.getObject(any())).thenAnswer((Answer<Struct>) invocation -> {
+            Struct structObj = Mockito.mock(Struct.class);
+            Object argument = invocation.getArguments()[0];
+            Mockito.when(structObj.getAttributes()).thenAnswer((Answer<Object[]>) invocation2 -> {
+                if (argument instanceof Integer) {
+                    int colIndex = (Integer) argument - 1;
+                    return (Object[]) rows[rowNumber.get()][colIndex];
+                }
+                else if (argument instanceof String) {
+                    int colIndex = Arrays.asList(columnNames).indexOf(argument);
+                    return (Object[]) rows[rowNumber.get()][colIndex];
+                }
+                else {
+                    throw new RuntimeException("Unexpected argument type " + argument.getClass());
+                }
+            });
+            return structObj;
+        });
+
+
         Mockito.when(resultSet.getDouble(any())).thenAnswer((Answer<Double>) invocation -> {
             Object argument = invocation.getArguments()[0];
             if (argument instanceof Integer) {