Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SQLServer glue get table api issue fix #2524

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
+30 −23
Open

SQLServer glue get table api issue fix #2524

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5546ae3
The SQLServer Glue Get Table API was not retrieving data as expected …
ritiktrianz Jan 17, 2025
ee2c964
The SQLServer Glue Get Table API was not retrieving data as expected …
ritiktrianz Jan 21, 2025
6c8e7dd
Merge branch 'master' into sqlserver_phase2_tablelist_issue
ritiktrianz Jan 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions ...sqlserver/src/main/java/com/amazonaws/athena/connectors/sqlserver/SqlServerConstants.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,5 +25,6 @@ public final class SqlServerConstants
public static final String DRIVER_CLASS = "com.microsoft.sqlserver.jdbc.SQLServerDriver";
public static final int DEFAULT_PORT = 1433;
public static final String SQLSERVER_QUOTE_CHARACTER = "\"";
static final String PARTITION_NUMBER = "partition_number";
private SqlServerConstants() {}
}
3 changes: 2 additions & 1 deletion ...ver/src/main/java/com/amazonaws/athena/connectors/sqlserver/SqlServerMetadataHandler.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,13 +84,14 @@
import java.util.Set;
import java.util.stream.Collectors;

import static com.amazonaws.athena.connectors.sqlserver.SqlServerConstants.PARTITION_NUMBER;

public class SqlServerMetadataHandler extends JdbcMetadataHandler
{
private static final Logger LOGGER = LoggerFactory.getLogger(SqlServerMetadataHandler.class);

static final Map<String, String> JDBC_PROPERTIES = ImmutableMap.of("databaseTerm", "SCHEMA");
static final String ALL_PARTITIONS = "0";
static final String PARTITION_NUMBER = "PARTITION_NUMBER";
static final String PARTITION_FUNCTION = "PARTITION_FUNCTION";
static final String PARTITIONING_COLUMN = "PARTITIONING_COLUMN";

Expand Down
4 changes: 3 additions & 1 deletion .../src/main/java/com/amazonaws/athena/connectors/sqlserver/SqlServerQueryStringBuilder.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,8 @@
import java.util.Collections;
import java.util.List;

import static com.amazonaws.athena.connectors.sqlserver.SqlServerConstants.PARTITION_NUMBER;

public class SqlServerQueryStringBuilder extends JdbcSplitQueryBuilder
{
private static final Logger LOGGER = LoggerFactory.getLogger(SqlServerQueryStringBuilder.class);
Expand Down Expand Up @@ -63,7 +65,7 @@ protected List<String> getPartitionWhereClauses(Split split)
{
String partitionFunction = split.getProperty(SqlServerMetadataHandler.PARTITION_FUNCTION);
String partitioningColumn = split.getProperty(SqlServerMetadataHandler.PARTITIONING_COLUMN);
String partitionNumber = split.getProperty(SqlServerMetadataHandler.PARTITION_NUMBER);
String partitionNumber = split.getProperty(PARTITION_NUMBER);

//example query: select * from MyPartitionTable where $PARTITION.myRangePF(col1) =2
LOGGER.debug("PARTITION_FUNCTION: {}", partitionFunction);
Expand Down
33 changes: 17 additions & 16 deletions ...src/test/java/com/amazonaws/athena/connectors/sqlserver/SqlServerMetadataHandlerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,14 +70,15 @@
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;

import static com.amazonaws.athena.connectors.sqlserver.SqlServerConstants.PARTITION_NUMBER;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.nullable;

public class SqlServerMetadataHandlerTest
extends TestBase
{
private static final Logger logger = LoggerFactory.getLogger(SqlServerMetadataHandlerTest.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", SqlServerConstants.NAME,
"sqlserver://jdbc:sqlserver://hostname;databaseName=fakedatabase");
private SqlServerMetadataHandler sqlServerMetadataHandler;
Expand Down Expand Up @@ -109,7 +110,7 @@ public void setup()
public void getPartitionSchema()
{
Assert.assertEquals(SchemaBuilder.newBuilder()
.addField(sqlServerMetadataHandler.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.sqlServerMetadataHandler.getPartitionSchema("testCatalogName"));
}

Expand All @@ -132,7 +133,7 @@ public void doGetTableLayout()
PreparedStatement preparedStatement = Mockito.mock(PreparedStatement.class);
Mockito.when(this.connection.prepareStatement(sqlServerMetadataHandler.GET_PARTITIONS_QUERY)).thenReturn(preparedStatement);

String[] columns = {sqlServerMetadataHandler.PARTITION_NUMBER};
String[] columns = {PARTITION_NUMBER};
int[] types = {Types.VARCHAR};
Object[][] values = {{"2"},{"3"}};
ResultSet resultSet = mockResultSet(columns, types, values, new AtomicInteger(-1));
Expand All @@ -150,10 +151,10 @@ public void doGetTableLayout()
for (int i = 0; i < getTableLayoutResponse.getPartitions().getRowCount(); i++) {
actualValues.add(BlockUtils.rowToString(getTableLayoutResponse.getPartitions(), i));
}
Assert.assertEquals(Arrays.asList("[PARTITION_NUMBER : 1:::pf:::pc]","[PARTITION_NUMBER : 2:::pf:::pc]","[PARTITION_NUMBER : 3:::pf:::pc]"), actualValues);
Assert.assertEquals(Arrays.asList("[partition_number : 1:::pf:::pc]","[partition_number : 2:::pf:::pc]","[partition_number : 3:::pf:::pc]"), actualValues);

SchemaBuilder expectedSchemaBuilder = SchemaBuilder.newBuilder();
expectedSchemaBuilder.addField(FieldBuilder.newBuilder(sqlServerMetadataHandler.PARTITION_NUMBER, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
expectedSchemaBuilder.addField(FieldBuilder.newBuilder(PARTITION_NUMBER, 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 @@ -176,7 +177,7 @@ public void doGetTableLayoutWithNoPartitions()
PreparedStatement preparedStatement = Mockito.mock(PreparedStatement.class);
Mockito.when(this.connection.prepareStatement(sqlServerMetadataHandler.GET_PARTITIONS_QUERY)).thenReturn(preparedStatement);

String[] columns = {sqlServerMetadataHandler.PARTITION_NUMBER};
String[] columns = {PARTITION_NUMBER};
int[] types = {Types.VARCHAR};
Object[][] values = {{}};
ResultSet resultSet = mockResultSet(columns, types, values, new AtomicInteger(-1));
Expand All @@ -196,10 +197,10 @@ public void doGetTableLayoutWithNoPartitions()
for (int i = 0; i < getTableLayoutResponse.getPartitions().getRowCount(); i++) {
actualValues.add(BlockUtils.rowToString(getTableLayoutResponse.getPartitions(), i));
}
Assert.assertEquals(Arrays.asList("[PARTITION_NUMBER : 0]"), actualValues);
Assert.assertEquals(Arrays.asList("[partition_number : 0]"), actualValues);

SchemaBuilder expectedSchemaBuilder = SchemaBuilder.newBuilder();
expectedSchemaBuilder.addField(FieldBuilder.newBuilder(sqlServerMetadataHandler.PARTITION_NUMBER, org.apache.arrow.vector.types.Types.MinorType.VARCHAR.getType()).build());
expectedSchemaBuilder.addField(FieldBuilder.newBuilder(PARTITION_NUMBER, 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 Down Expand Up @@ -248,7 +249,7 @@ public void doGetSplits()
PreparedStatement preparedStatement = Mockito.mock(PreparedStatement.class);
Mockito.when(this.connection.prepareStatement(sqlServerMetadataHandler.GET_PARTITIONS_QUERY)).thenReturn(preparedStatement);

String[] columns = {sqlServerMetadataHandler.PARTITION_NUMBER};
String[] columns = {PARTITION_NUMBER};
int[] types = {Types.INTEGER};
Object[][] values = {{2}, {3}};
ResultSet resultSet = mockResultSet(columns, types, values, new AtomicInteger(-1));
Expand All @@ -273,15 +274,15 @@ public void doGetSplits()

Set<Map<String, String>> expectedSplits = com.google.common.collect.ImmutableSet.of(
com.google.common.collect.ImmutableMap.of(
sqlServerMetadataHandler.PARTITION_NUMBER, "1",
PARTITION_NUMBER, "1",
"PARTITIONING_COLUMN", "pc",
"PARTITION_FUNCTION", "pf"),
com.google.common.collect.ImmutableMap.of(
sqlServerMetadataHandler.PARTITION_NUMBER, "2",
PARTITION_NUMBER, "2",
"PARTITIONING_COLUMN", "pc",
"PARTITION_FUNCTION", "pf"),
com.google.common.collect.ImmutableMap.of(
sqlServerMetadataHandler.PARTITION_NUMBER, "3",
PARTITION_NUMBER, "3",
"PARTITIONING_COLUMN", "pc",
"PARTITION_FUNCTION", "pf"));
Assert.assertEquals(expectedSplits.size(), getSplitsResponse.getSplits().size());
Expand Down Expand Up @@ -310,7 +311,7 @@ public void doGetSplitsWithNoPartition()
PreparedStatement preparedStatement = Mockito.mock(PreparedStatement.class);
Mockito.when(this.connection.prepareStatement(sqlServerMetadataHandler.GET_PARTITIONS_QUERY)).thenReturn(preparedStatement);

String[] columns = {sqlServerMetadataHandler.PARTITION_NUMBER};
String[] columns = {PARTITION_NUMBER};
int[] types = {Types.INTEGER};
Object[][] values = {{}};
ResultSet resultSet = mockResultSet(columns, types, values, new AtomicInteger(-1));
Expand All @@ -329,7 +330,7 @@ public void doGetSplitsWithNoPartition()
GetSplitsResponse getSplitsResponse = this.sqlServerMetadataHandler.doGetSplits(splitBlockAllocator, getSplitsRequest);

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

String[] columns = {sqlServerMetadataHandler.PARTITION_NUMBER};
String[] columns = {PARTITION_NUMBER};
int[] types = {Types.INTEGER};
Object[][] values = {{2},{3}};
ResultSet resultSet = mockResultSet(columns, types, values, new AtomicInteger(-1));
Expand All @@ -375,7 +376,7 @@ public void doGetSplitsContinuation()

Set<Map<String, String>> expectedSplits = com.google.common.collect.ImmutableSet.of(
com.google.common.collect.ImmutableMap.of(
sqlServerMetadataHandler.PARTITION_NUMBER, "3",
PARTITION_NUMBER, "3",
"PARTITIONING_COLUMN", "pc",
"PARTITION_FUNCTION", "pf"));
Set<Map<String, String>> actualSplits = getSplitsResponse.getSplits().stream().map(Split::getProperties).collect(Collectors.toSet());
Expand Down
9 changes: 5 additions & 4 deletions .../test/java/com/amazonaws/athena/connectors/sqlserver/SqlServerQueryStringBuilderTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import java.util.Collections;
import java.util.List;

import static com.amazonaws.athena.connectors.sqlserver.SqlServerConstants.PARTITION_NUMBER;
import static com.amazonaws.athena.connectors.sqlserver.SqlServerConstants.SQLSERVER_QUOTE_CHARACTER;

public class SqlServerQueryStringBuilderTest
Expand Down Expand Up @@ -61,7 +62,7 @@ public void testGetPartitionWhereClauses()
Split split1 = Mockito.mock(Split.class);
Mockito.when(split1.getProperty(SqlServerMetadataHandler.PARTITION_FUNCTION)).thenReturn("pf");
Mockito.when(split1.getProperty(SqlServerMetadataHandler.PARTITIONING_COLUMN)).thenReturn("col");
Mockito.when(split1.getProperty(SqlServerMetadataHandler.PARTITION_NUMBER)).thenReturn("1");
Mockito.when(split1.getProperty(PARTITION_NUMBER)).thenReturn("1");
Assert.assertEquals(Collections.singletonList(" $PARTITION.pf(col) = 1"), builder.getPartitionWhereClauses(split1));

}
Expand All @@ -73,7 +74,7 @@ public void getPartitionWhereClausesWhenPartitionFunctionIsNull()
Split split = Mockito.mock(Split.class);
Mockito.when(split.getProperty(SqlServerMetadataHandler.PARTITION_FUNCTION)).thenReturn(null);
Mockito.when(split.getProperty(SqlServerMetadataHandler.PARTITIONING_COLUMN)).thenReturn("col");
Mockito.when(split.getProperty(SqlServerMetadataHandler.PARTITION_NUMBER)).thenReturn("74");
Mockito.when(split.getProperty(PARTITION_NUMBER)).thenReturn("74");
List<String> result = builder.getPartitionWhereClauses(split);
// Should return empty list if partition function is null.
Assert.assertEquals(0, result.size());
Expand All @@ -86,7 +87,7 @@ public void getPartitionWhereClausesWhenPartitioningColumnIsNull()
Split split = Mockito.mock(Split.class);
Mockito.when(split.getProperty(SqlServerMetadataHandler.PARTITION_FUNCTION)).thenReturn("pf");
Mockito.when(split.getProperty(SqlServerMetadataHandler.PARTITIONING_COLUMN)).thenReturn(null);
Mockito.when(split.getProperty(SqlServerMetadataHandler.PARTITION_NUMBER)).thenReturn("74");
Mockito.when(split.getProperty(PARTITION_NUMBER)).thenReturn("74");
List<String> result = builder.getPartitionWhereClauses(split);
// Should return empty list if partitioning column is null.
Assert.assertEquals(0, result.size());
Expand All @@ -99,7 +100,7 @@ public void getPartitionWhereClausesWhenPartitionFunctionAndColumnAreNull()
Split split = Mockito.mock(Split.class);
Mockito.when(split.getProperty(SqlServerMetadataHandler.PARTITION_FUNCTION)).thenReturn(null);
Mockito.when(split.getProperty(SqlServerMetadataHandler.PARTITIONING_COLUMN)).thenReturn(null);
Mockito.when(split.getProperty(SqlServerMetadataHandler.PARTITION_NUMBER)).thenReturn("74");
Mockito.when(split.getProperty(PARTITION_NUMBER)).thenReturn("74");
List<String> result = builder.getPartitionWhereClauses(split);
// Should return empty list if both partition function and partitioning column are null.
Assert.assertEquals(0, result.size());
Expand Down
3 changes: 2 additions & 1 deletion ...r/src/test/java/com/amazonaws/athena/connectors/sqlserver/SqlServerRecordHandlerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@
import java.sql.SQLException;
import java.util.Collections;

import static com.amazonaws.athena.connectors.sqlserver.SqlServerConstants.PARTITION_NUMBER;
import static com.amazonaws.athena.connectors.sqlserver.SqlServerConstants.SQLSERVER_QUOTE_CHARACTER;
import static org.mockito.ArgumentMatchers.nullable;

Expand Down Expand Up @@ -104,7 +105,7 @@ public void buildSplitSqlNew()
Split split = Mockito.mock(Split.class);
Mockito.when(split.getProperty(SqlServerMetadataHandler.PARTITION_FUNCTION)).thenReturn("pf");
Mockito.when(split.getProperty(SqlServerMetadataHandler.PARTITIONING_COLUMN)).thenReturn("testCol1");
Mockito.when(split.getProperty(SqlServerMetadataHandler.PARTITION_NUMBER)).thenReturn("1");
Mockito.when(split.getProperty(PARTITION_NUMBER)).thenReturn("1");

ValueSet valueSet = getSingleValueSet("varcharTest");
Constraints constraints = Mockito.mock(Constraints.class);
Expand Down