Skip to content

Commit

Permalink
Saphana glue get table api issue fix (#2505)
Browse files Browse the repository at this point in the history 
Co-authored-by: akshay.kachore <akshay.kachore@trianz.com>
  • Loading branch information
@VenkatasivareddyTR @Trianz-Akshay
VenkatasivareddyTR and Trianz-Akshay authored Jan 15, 2025
1 parent a1a4242 commit 01bd504
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 17 deletions.
4 changes: 2 additions & 2 deletions athena-saphana/src/main/java/com/amazonaws/athena/connectors/saphana/SaphanaConstants.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,15 +32,15 @@ public final class SaphanaConstants
public static final int SAPHANA_DEFAULT_PORT = 1025;
static final Map<String, String> JDBC_PROPERTIES = ImmutableMap.of("databaseTerm", "SCHEMA");
static final String ALL_PARTITIONS = "0";
static final String BLOCK_PARTITION_COLUMN_NAME = "PART_ID";
static final String BLOCK_PARTITION_COLUMN_NAME = "part_id";
static final String COLUMN_NAME = "COLUMN_NAME";
static final String CASE_UPPER = "upper";
static final String CASE_LOWER = "lower";
public static final String SAPHANA_QUOTE_CHARACTER = "\"";
/**
* partition query for saphana
*/
static final String GET_PARTITIONS_QUERY = "SELECT DISTINCT PART_ID FROM SYS.TABLE_PARTITIONS " +
static final String GET_PARTITIONS_QUERY = "SELECT DISTINCT PART_ID as \"part_id\" FROM SYS.TABLE_PARTITIONS " +
"WHERE TABLE_NAME = ? AND SCHEMA_NAME = ? AND PART_ID IS NOT NULL";

/**
Expand Down
31 changes: 16 additions & 15 deletions ...ana/src/test/java/com/amazonaws/athena/connectors/saphana/SaphanaMetadataHandlerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@
import software.amazon.awssdk.services.secretsmanager.model.GetSecretValueRequest;
import software.amazon.awssdk.services.secretsmanager.model.GetSecretValueResponse;

import static com.amazonaws.athena.connectors.saphana.SaphanaConstants.BLOCK_PARTITION_COLUMN_NAME;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.nullable;

Expand All @@ -68,7 +69,7 @@ public class SaphanaMetadataHandlerTest
private SecretsManagerClient secretsManager;
private AthenaClient athena;
private BlockAllocator blockAllocator;
private static final Schema PARTITION_SCHEMA = SchemaBuilder.newBuilder().addField("PART_ID", org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build();
private static final Schema PARTITION_SCHEMA = SchemaBuilder.newBuilder().addField(BLOCK_PARTITION_COLUMN_NAME, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build();


@Before
Expand All @@ -90,7 +91,7 @@ public void setup()
public void getPartitionSchema()
{
Assert.assertEquals(SchemaBuilder.newBuilder()
.addField(SaphanaConstants.BLOCK_PARTITION_COLUMN_NAME, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build(),
.addField(BLOCK_PARTITION_COLUMN_NAME, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build(),
this.saphanaMetadataHandler.getPartitionSchema("testCatalogName"));
}

Expand All @@ -102,14 +103,14 @@ public void doGetTableLayout()
Constraints constraints = Mockito.mock(Constraints.class);
TableName tableName = new TableName("testSchema", "testTable");
Schema partitionSchema = this.saphanaMetadataHandler.getPartitionSchema("testCatalogName");
Set<String> partitionCols = new HashSet<>(Arrays.asList("PART_ID")); //partitionSchema.getFields().stream().map(Field::getName).collect(Collectors.toSet());
Set<String> partitionCols = new HashSet<>(Arrays.asList(BLOCK_PARTITION_COLUMN_NAME)); //partitionSchema.getFields().stream().map(Field::getName).collect(Collectors.toSet());
GetTableLayoutRequest getTableLayoutRequest = new GetTableLayoutRequest(this.federatedIdentity, "testQueryId", "testCatalogName", tableName, constraints, partitionSchema, partitionCols);

PreparedStatement preparedStatement = Mockito.mock(PreparedStatement.class);

Mockito.when(this.connection.prepareStatement(SaphanaConstants.GET_PARTITIONS_QUERY)).thenReturn(preparedStatement);

String[] columns = {"PART_ID"};
String[] columns = {BLOCK_PARTITION_COLUMN_NAME};
int[] types = {Types.VARCHAR};
Object[][] values = {{"p0"}, {"p1"}};
ResultSet resultSet = mockResultSet(columns, types, values, new AtomicInteger(-1));
Expand All @@ -123,10 +124,10 @@ public void doGetTableLayout()
for (int i = 0; i < getTableLayoutResponse.getPartitions().getRowCount(); i++) {
expectedValues.add(BlockUtils.rowToString(getTableLayoutResponse.getPartitions(), i));
}
Assert.assertEquals(expectedValues, Arrays.asList("[PART_ID : p0]", "[PART_ID : p1]"));
Assert.assertEquals(expectedValues, Arrays.asList("[part_id : p0]", "[part_id : p1]"));

SchemaBuilder expectedSchemaBuilder = SchemaBuilder.newBuilder();
expectedSchemaBuilder.addField(FieldBuilder.newBuilder(SaphanaConstants.BLOCK_PARTITION_COLUMN_NAME, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
expectedSchemaBuilder.addField(FieldBuilder.newBuilder(BLOCK_PARTITION_COLUMN_NAME, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
Schema expectedSchema = expectedSchemaBuilder.build();
Assert.assertEquals(expectedSchema, getTableLayoutResponse.getPartitions().getSchema());
Assert.assertEquals(tableName, getTableLayoutResponse.getTableName());
Expand All @@ -149,7 +150,7 @@ public void doGetTableLayoutWithNoPartitions()
PreparedStatement preparedStatement = Mockito.mock(PreparedStatement.class);
Mockito.when(this.connection.prepareStatement(SaphanaConstants.GET_PARTITIONS_QUERY)).thenReturn(preparedStatement);

String[] columns = {"PART_ID"};
String[] columns = {BLOCK_PARTITION_COLUMN_NAME};
int[] types = {Types.VARCHAR};
Object[][] values = {{}};
ResultSet resultSet = mockResultSet(columns, types, values, new AtomicInteger(-1));
Expand All @@ -165,10 +166,10 @@ public void doGetTableLayoutWithNoPartitions()
for (int i = 0; i < getTableLayoutResponse.getPartitions().getRowCount(); i++) {
expectedValues.add(BlockUtils.rowToString(getTableLayoutResponse.getPartitions(), i));
}
Assert.assertEquals(expectedValues, Collections.singletonList("[PART_ID : 0]"));
Assert.assertEquals(expectedValues, Collections.singletonList("[part_id : 0]"));

SchemaBuilder expectedSchemaBuilder = SchemaBuilder.newBuilder();
expectedSchemaBuilder.addField(FieldBuilder.newBuilder(SaphanaConstants.BLOCK_PARTITION_COLUMN_NAME, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
expectedSchemaBuilder.addField(FieldBuilder.newBuilder(BLOCK_PARTITION_COLUMN_NAME, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
Assert.assertEquals(tableName, getTableLayoutResponse.getTableName());

Mockito.verify(preparedStatement, Mockito.times(1)).setString(1, tableName.getTableName());
Expand Down Expand Up @@ -203,7 +204,7 @@ public void doGetSplits()

PreparedStatement preparedStatement = Mockito.mock(PreparedStatement.class);

String[] columns = {SaphanaConstants.BLOCK_PARTITION_COLUMN_NAME};
String[] columns = {BLOCK_PARTITION_COLUMN_NAME};
int[] types = {Types.VARCHAR};
Object[][] values = {{"p0"}, {"p1"}};
ResultSet resultSet = mockResultSet(columns, types, values, new AtomicInteger(-1));
Expand All @@ -224,8 +225,8 @@ public void doGetSplits()
GetSplitsResponse getSplitsResponse = this.saphanaMetadataHandler.doGetSplits(splitBlockAllocator, getSplitsRequest);

Set<Map<String, String>> expectedSplits = new HashSet<>();
expectedSplits.add(Collections.singletonMap(SaphanaConstants.BLOCK_PARTITION_COLUMN_NAME, "p0"));
expectedSplits.add(Collections.singletonMap(SaphanaConstants.BLOCK_PARTITION_COLUMN_NAME, "p1"));
expectedSplits.add(Collections.singletonMap(BLOCK_PARTITION_COLUMN_NAME, "p0"));
expectedSplits.add(Collections.singletonMap(BLOCK_PARTITION_COLUMN_NAME, "p1"));
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);
Expand All @@ -245,7 +246,7 @@ public void doGetSplitsContinuation()
PreparedStatement preparedStatement = Mockito.mock(PreparedStatement.class);
Mockito.when(this.connection.prepareStatement(SaphanaConstants.GET_PARTITIONS_QUERY)).thenReturn(preparedStatement);

String[] columns = {"PART_ID"};
String[] columns = {BLOCK_PARTITION_COLUMN_NAME};
int[] types = {Types.VARCHAR};
Object[][] values = {{"p0"}, {"p1"}};
ResultSet resultSet = mockResultSet(columns, types, values, new AtomicInteger(-1));
Expand All @@ -260,7 +261,7 @@ public void doGetSplitsContinuation()
GetSplitsResponse getSplitsResponse = this.saphanaMetadataHandler.doGetSplits(splitBlockAllocator, getSplitsRequest);

Set<Map<String, String>> expectedSplits = new HashSet<>();
expectedSplits.add(Collections.singletonMap("PART_ID", "p1"));
expectedSplits.add(Collections.singletonMap(BLOCK_PARTITION_COLUMN_NAME, "p1"));
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);
Expand Down Expand Up @@ -339,7 +340,7 @@ public void doGetSplitsForView()
GetSplitsResponse getSplitsResponse = this.saphanaMetadataHandler.doGetSplits(splitBlockAllocator, getSplitsRequest);

Set<Map<String, String>> expectedSplits = new HashSet<>();
expectedSplits.add(Collections.singletonMap(SaphanaConstants.BLOCK_PARTITION_COLUMN_NAME, "0"));
expectedSplits.add(Collections.singletonMap(BLOCK_PARTITION_COLUMN_NAME, "0"));

Assert.assertEquals(expectedSplits.size(), getSplitsResponse.getSplits().size());

Expand Down

0 comments on commit 01bd504

Please sign in to comment.