DB2/DB2as400 glue get table api issue fix (#2506)

athena-db2-as400/src/main/java/com/amazonaws/athena/connectors/db2as400/Db2As400MetadataHandler.java

@@ -77,7 +77,7 @@
 public class Db2As400MetadataHandler extends JdbcMetadataHandler
 {
     private static final Logger LOGGER = LoggerFactory.getLogger(Db2As400MetadataHandler.class);
-    static final String PARTITION_NUMBER = "PARTITION_NUMBER";
+    static final String PARTITION_NUMBER = "partition_number";
     static final String PARTITIONING_COLUMN = "PARTITIONING_COLUMN";
     /**
      * DB2 has max number of partition 32,000

athena-db2-as400/src/test/java/com/amazonaws/athena/connectors/db2as400/Db2As400MetadataHandlerTest.java

@@ -69,11 +69,12 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 
+import static com.amazonaws.athena.connectors.db2as400.Db2As400MetadataHandler.PARTITION_NUMBER;
 import static org.mockito.ArgumentMatchers.nullable;
 
 public class Db2As400MetadataHandlerTest extends TestBase {
     private static final Logger logger = LoggerFactory.getLogger(Db2As400MetadataHandlerTest.class);
-    private static final Schema PARTITION_SCHEMA = SchemaBuilder.newBuilder().addField("PARTITION_NUMBER", org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build();
+    private static final Schema PARTITION_SCHEMA = SchemaBuilder.newBuilder().addField(PARTITION_NUMBER, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build();
     private DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", Db2As400Constants.NAME,
             "db2as400://jdbc:as400://testhost;user=dummy;password=dummy;");
     private Db2As400MetadataHandler db2As400MetadataHandler;
@@ -103,7 +104,7 @@ public void setup() throws Exception {
     public void getPartitionSchema()
     {
         Assert.assertEquals(SchemaBuilder.newBuilder()
-                        .addField(Db2As400MetadataHandler.PARTITION_NUMBER, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build(),
+                        .addField(PARTITION_NUMBER, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build(),
                 this.db2As400MetadataHandler.getPartitionSchema("testCatalogName"));
     }
 
@@ -130,7 +131,7 @@ public void doGetSplitsWithNoPartition()
         GetSplitsResponse getSplitsResponse = this.db2As400MetadataHandler.doGetSplits(splitBlockAllocator, getSplitsRequest);
 
         Set<Map<String, String>> expectedSplits = new HashSet<>();
-        expectedSplits.add(Collections.singletonMap(db2As400MetadataHandler.PARTITION_NUMBER, "0"));
+        expectedSplits.add(Collections.singletonMap(PARTITION_NUMBER, "0"));
         Assert.assertEquals(expectedSplits.size(), getSplitsResponse.getSplits().size());
         Set<Map<String, String>> actualSplits = getSplitsResponse.getSplits().stream().map(Split::getProperties).collect(Collectors.toSet());
         Assert.assertEquals(expectedSplits, actualSplits);
@@ -167,13 +168,13 @@ public void doGetSplits()
 
         Set<Map<String, String>> expectedSplits = com.google.common.collect.ImmutableSet.of(
             com.google.common.collect.ImmutableMap.of(
-                db2As400MetadataHandler.PARTITION_NUMBER, "0",
+                PARTITION_NUMBER, "0",
                 db2As400MetadataHandler.PARTITIONING_COLUMN, "PC"),
             com.google.common.collect.ImmutableMap.of(
-                db2As400MetadataHandler.PARTITION_NUMBER, "1",
+                PARTITION_NUMBER, "1",
                 db2As400MetadataHandler.PARTITIONING_COLUMN, "PC"),
             com.google.common.collect.ImmutableMap.of(
-                db2As400MetadataHandler.PARTITION_NUMBER, "2",
+                PARTITION_NUMBER, "2",
                 db2As400MetadataHandler.PARTITIONING_COLUMN, "PC"));
 
         Assert.assertEquals(expectedSplits.size(), getSplitsResponse.getSplits().size());

athena-db2-as400/src/test/java/com/amazonaws/athena/connectors/db2as400/Db2As400QueryStringBuilderTest.java

@@ -45,7 +45,7 @@ public void testGetPartitionWhereClauses()
     {
         Db2As400QueryStringBuilder builder = new Db2As400QueryStringBuilder("'");
         Split split = Mockito.mock(Split.class);
-        Mockito.when(split.getProperty(Mockito.eq("PARTITION_NUMBER"))).thenReturn("0");
+        Mockito.when(split.getProperty(Mockito.eq("partition_number"))).thenReturn("0");
         Mockito.when(split.getProperty(Mockito.eq("PARTITIONING_COLUMN"))).thenReturn("PC");
         Assert.assertEquals(Arrays.asList(" DATAPARTITIONNUM(PC) = 0"), builder.getPartitionWhereClauses(split));
     }

athena-db2/src/main/java/com/amazonaws/athena/connectors/db2/Db2Constants.java

@@ -25,7 +25,7 @@ public class Db2Constants
     public static final String DRIVER_CLASS = "com.ibm.db2.jcc.DB2Driver";
     public static final int DEFAULT_PORT = 50001;
     public static final String QUOTE_CHARACTER = "\"";
-
+    static final String PARTITION_NUMBER = "partition_number";
     public static final String QRY_TO_LIST_SCHEMAS = "select schemaname as name " +
             "from syscat.schemata " +
             "where schemaname not like 'SYS%' " +

athena-db2/src/main/java/com/amazonaws/athena/connectors/db2/Db2MetadataHandler.java

@@ -82,11 +82,11 @@
 import java.util.stream.Collectors;
 
 import static com.amazonaws.athena.connector.lambda.domain.predicate.functions.StandardFunctions.NULLIF_FUNCTION_NAME;
+import static com.amazonaws.athena.connectors.db2.Db2Constants.PARTITION_NUMBER;
 
 public class Db2MetadataHandler extends JdbcMetadataHandler
 {
     private static final Logger LOGGER = LoggerFactory.getLogger(Db2MetadataHandler.class);
-    static final String PARTITION_NUMBER = "PARTITION_NUMBER";
     static final String PARTITIONING_COLUMN = "PARTITIONING_COLUMN";
     /**
      * DB2 has max number of partition 32,000

athena-db2/src/main/java/com/amazonaws/athena/connectors/db2/Db2QueryStringBuilder.java

@@ -29,6 +29,8 @@
 import java.util.Collections;
 import java.util.List;
 
+import static com.amazonaws.athena.connectors.db2.Db2Constants.PARTITION_NUMBER;
+
 public class Db2QueryStringBuilder extends JdbcSplitQueryBuilder
 {
     private static final Logger LOGGER = LoggerFactory.getLogger(Db2QueryStringBuilder.class);
@@ -74,7 +76,7 @@ protected List<String> getPartitionWhereClauses(Split split)
         if (column != null) {
             LOGGER.debug("Fetching data using Partition");
             //example query: select * from EMP_TABLE WHERE DATAPARTITIONNUM(EMP_NO) = 0
-            return Collections.singletonList(" DATAPARTITIONNUM(" + column + ") = " + split.getProperty(Db2MetadataHandler.PARTITION_NUMBER));
+            return Collections.singletonList(" DATAPARTITIONNUM(" + column + ") = " + split.getProperty(PARTITION_NUMBER));
         }
         else {
             LOGGER.debug("Fetching data without Partition");

athena-db2/src/test/java/com/amazonaws/athena/connectors/db2/Db2MetadataHandlerTest.java

@@ -69,12 +69,13 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 
+import static com.amazonaws.athena.connectors.db2.Db2Constants.PARTITION_NUMBER;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.nullable;
 
 public class Db2MetadataHandlerTest extends TestBase {
     private static final Logger logger = LoggerFactory.getLogger(Db2MetadataHandlerTest.class);
-    private static final Schema PARTITION_SCHEMA = SchemaBuilder.newBuilder().addField("PARTITION_NUMBER", org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build();
+    private static final Schema PARTITION_SCHEMA = SchemaBuilder.newBuilder().addField(PARTITION_NUMBER, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build();
     private DatabaseConnectionConfig databaseConnectionConfig = new DatabaseConnectionConfig("testCatalog", Db2Constants.NAME,
             "dbtwo://jdbc:db2://hostname:50001/dummydatabase:user=dummyuser;password=dummypwd");
     private Db2MetadataHandler db2MetadataHandler;
@@ -104,7 +105,7 @@ public void setup() throws Exception {
     public void getPartitionSchema()
     {
         Assert.assertEquals(SchemaBuilder.newBuilder()
-                        .addField(Db2MetadataHandler.PARTITION_NUMBER, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build(),
+                        .addField(PARTITION_NUMBER, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build(),
                 this.db2MetadataHandler.getPartitionSchema("testCatalogName"));
     }
 
@@ -131,7 +132,7 @@ public void doGetSplitsWithNoPartition()
         GetSplitsResponse getSplitsResponse = this.db2MetadataHandler.doGetSplits(splitBlockAllocator, getSplitsRequest);
 
         Set<Map<String, String>> expectedSplits = new HashSet<>();
-        expectedSplits.add(Collections.singletonMap(db2MetadataHandler.PARTITION_NUMBER, "0"));
+        expectedSplits.add(Collections.singletonMap(PARTITION_NUMBER, "0"));
         Assert.assertEquals(expectedSplits.size(), getSplitsResponse.getSplits().size());
         Set<Map<String, String>> actualSplits = getSplitsResponse.getSplits().stream().map(Split::getProperties).collect(Collectors.toSet());
         Assert.assertEquals(expectedSplits, actualSplits);
@@ -168,13 +169,13 @@ public void doGetSplits()
 
         Set<Map<String, String>> expectedSplits = com.google.common.collect.ImmutableSet.of(
             com.google.common.collect.ImmutableMap.of(
-                db2MetadataHandler.PARTITION_NUMBER, "0",
+                PARTITION_NUMBER, "0",
                 db2MetadataHandler.PARTITIONING_COLUMN, "PC"),
             com.google.common.collect.ImmutableMap.of(
-                db2MetadataHandler.PARTITION_NUMBER, "1",
+                PARTITION_NUMBER, "1",
                 db2MetadataHandler.PARTITIONING_COLUMN, "PC"),
             com.google.common.collect.ImmutableMap.of(
-                db2MetadataHandler.PARTITION_NUMBER, "2",
+                PARTITION_NUMBER, "2",
                 db2MetadataHandler.PARTITIONING_COLUMN, "PC"));
 
         Assert.assertEquals(expectedSplits.size(), getSplitsResponse.getSplits().size());

athena-db2/src/test/java/com/amazonaws/athena/connectors/db2/Db2QueryStringBuilderTest.java

@@ -47,7 +47,7 @@ public void testGetPartitionWhereClauses()
     {
         Db2QueryStringBuilder builder =  new Db2QueryStringBuilder(QUOTE_CHARACTER, new Db2FederationExpressionParser(QUOTE_CHARACTER));
         Split split = Mockito.mock(Split.class);
-        Mockito.when(split.getProperty(Mockito.eq("PARTITION_NUMBER"))).thenReturn("0");
+        Mockito.when(split.getProperty(Mockito.eq("partition_number"))).thenReturn("0");
         Mockito.when(split.getProperty(Mockito.eq("PARTITIONING_COLUMN"))).thenReturn("PC");
         Assert.assertEquals(Arrays.asList(" DATAPARTITIONNUM(PC) = 0"), builder.getPartitionWhereClauses(split));
     }

athena-db2/src/test/java/com/amazonaws/athena/connectors/db2/Db2RecordHandlerTest.java

@@ -47,6 +47,7 @@
 import java.sql.SQLException;
 import java.util.Collections;
 
+import static com.amazonaws.athena.connectors.db2.Db2Constants.PARTITION_NUMBER;
 import static com.amazonaws.athena.connectors.db2.Db2Constants.QUOTE_CHARACTER;
 import static org.mockito.ArgumentMatchers.nullable;
 
@@ -97,7 +98,7 @@ public void buildSplitSqlNew()
         Schema schema = schemaBuilder.build();
 
         Split split = Mockito.mock(Split.class);
-        Mockito.when(split.getProperty(Db2MetadataHandler.PARTITION_NUMBER)).thenReturn("0");
+        Mockito.when(split.getProperty(PARTITION_NUMBER)).thenReturn("0");
 
         ValueSet valueSet = getSingleValueSet("varcharTest");
         Constraints constraints = Mockito.mock(Constraints.class);