Skip to content

Commit 100b1dd

Browse files
Jithendar12Jithendar12
committed
Add case-insensitive search in DB2 Connector
1 parent 6d6e629 commit 100b1dd

File tree

5 files changed

+197
-11
lines changed

5 files changed

+197
-11
lines changed

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

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
import com.amazonaws.athena.connector.lambda.metadata.optimizations.pushdown.FilterPushdownSubType;
4646
import com.amazonaws.athena.connector.lambda.metadata.optimizations.pushdown.LimitPushdownSubType;
4747
import com.amazonaws.athena.connector.lambda.metadata.optimizations.pushdown.TopNPushdownSubType;
48+
import com.amazonaws.athena.connectors.db2.resolver.Db2JDBCCaseResolver;
4849
import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig;
4950
import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionInfo;
5051
import com.amazonaws.athena.connectors.jdbc.connection.GenericJdbcConnectionFactory;
@@ -53,6 +54,7 @@
5354
import com.amazonaws.athena.connectors.jdbc.manager.JdbcArrowTypeConverter;
5455
import com.amazonaws.athena.connectors.jdbc.manager.JdbcMetadataHandler;
5556
import com.amazonaws.athena.connectors.jdbc.manager.PreparedStatementBuilder;
57+
import com.amazonaws.athena.connectors.jdbc.resolver.JDBCCaseResolver;
5658
import com.google.common.annotations.VisibleForTesting;
5759
import com.google.common.collect.ImmutableMap;
5860
import com.google.common.collect.ImmutableSet;
@@ -121,18 +123,19 @@ public Db2MetadataHandler(
121123
JdbcConnectionFactory jdbcConnectionFactory,
122124
java.util.Map<String, String> configOptions)
123125
{
124-
super(databaseConnectionConfig, jdbcConnectionFactory, configOptions);
126+
super(databaseConnectionConfig, jdbcConnectionFactory, configOptions, new Db2JDBCCaseResolver(Db2Constants.NAME));
125127
}
126128

127129
@VisibleForTesting
128130
protected Db2MetadataHandler(
129-
DatabaseConnectionConfig databaseConnectionConfig,
130-
SecretsManagerClient secretsManager,
131-
AthenaClient athena,
132-
JdbcConnectionFactory jdbcConnectionFactory,
133-
java.util.Map<String, String> configOptions)
131+
DatabaseConnectionConfig databaseConnectionConfig,
132+
SecretsManagerClient secretsManager,
133+
AthenaClient athena,
134+
JdbcConnectionFactory jdbcConnectionFactory,
135+
java.util.Map<String, String> configOptions,
136+
JDBCCaseResolver caseResolver)
134137
{
135-
super(databaseConnectionConfig, secretsManager, athena, jdbcConnectionFactory, configOptions);
138+
super(databaseConnectionConfig, secretsManager, athena, jdbcConnectionFactory, configOptions, caseResolver);
136139
}
137140

138141
/**
@@ -163,8 +166,9 @@ public ListTablesResponse doListTables(final BlockAllocator blockAllocator, fina
163166
{
164167
try (Connection connection = getJdbcConnectionFactory().getConnection(getCredentialProvider())) {
165168
LOGGER.info("{}: List table names for Catalog {}, Schema {}", listTablesRequest.getQueryId(), listTablesRequest.getCatalogName(), listTablesRequest.getSchemaName());
166-
List<String> tableNames = getTableList(connection, Db2Constants.QRY_TO_LIST_TABLES_AND_VIEWS, listTablesRequest.getSchemaName());
167-
List<TableName> tables = tableNames.stream().map(tableName -> new TableName(listTablesRequest.getSchemaName(), tableName)).collect(Collectors.toList());
169+
String adjustedSchemaName = caseResolver.getAdjustedSchemaNameString(connection, listTablesRequest.getSchemaName(), configOptions);
170+
List<String> tableNames = getTableList(connection, Db2Constants.QRY_TO_LIST_TABLES_AND_VIEWS, adjustedSchemaName);
171+
List<TableName> tables = tableNames.stream().map(tableName -> new TableName(adjustedSchemaName, tableName)).collect(Collectors.toList());
168172
return new ListTablesResponse(listTablesRequest.getCatalogName(), tables, null);
169173
}
170174
}

‎athena-db2/src/main/java/com/amazonaws/athena/connectors/db2/Db2MuxMetadataHandler.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
*/
2020
package com.amazonaws.athena.connectors.db2;
2121

22+
import com.amazonaws.athena.connectors.db2.resolver.Db2JDBCCaseResolver;
2223
import com.amazonaws.athena.connectors.jdbc.MultiplexingJdbcMetadataHandler;
2324
import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig;
2425
import com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory;
@@ -58,6 +59,6 @@ public Db2MuxMetadataHandler(java.util.Map<String, String> configOptions)
5859
protected Db2MuxMetadataHandler(SecretsManagerClient secretsManager, AthenaClient athena, JdbcConnectionFactory jdbcConnectionFactory,
5960
Map<String, JdbcMetadataHandler> metadataHandlerMap, DatabaseConnectionConfig databaseConnectionConfig, java.util.Map<String, String> configOptions)
6061
{
61-
super(secretsManager, athena, jdbcConnectionFactory, metadataHandlerMap, databaseConnectionConfig, configOptions);
62+
super(secretsManager, athena, jdbcConnectionFactory, metadataHandlerMap, databaseConnectionConfig, configOptions, new Db2JDBCCaseResolver(Db2Constants.NAME));
6263
}
6364
}

‎athena-db2/src/main/java/com/amazonaws/athena/connectors/db2/resolver/Db2JDBCCaseResolver.java‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
/*-
2+
* #%L
3+
* athena-db2
4+
* %%
5+
* Copyright (C) 2019 - 2025 Amazon Web Services
6+
* %%
7+
* Licensed under the Apache License, Version 2.0 (the "License");
8+
* you may not use this file except in compliance with the License.
9+
* You may obtain a copy of the License at
10+
*
11+
* http://www.apache.org/licenses/LICENSE-2.0
12+
*
13+
* Unless required by applicable law or agreed to in writing, software
14+
* distributed under the License is distributed on an "AS IS" BASIS,
15+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16+
* See the License for the specific language governing permissions and
17+
* limitations under the License.
18+
* #L%
19+
*/
20+
package com.amazonaws.athena.connectors.db2.resolver;
21+
22+
import com.amazonaws.athena.connectors.jdbc.resolver.DefaultJDBCCaseResolver;
23+
24+
import java.util.List;
25+
26+
public class Db2JDBCCaseResolver
27+
extends DefaultJDBCCaseResolver
28+
{
29+
private static final String TABLE_NAME_QUERY_TEMPLATE = "SELECT TABNAME FROM SYSCAT.TABLES WHERE TABSCHEMA = ? AND LOWER(TABNAME) = ?";
30+
private static final String SCHEMA_NAME_QUERY_TEMPLATE = "SELECT SCHEMANAME FROM SYSCAT.SCHEMATA WHERE LOWER(SCHEMANAME) = ?";
31+
32+
private static final String SCHEMA_NAME_COLUMN_KEY = "SCHEMANAME";
33+
private static final String TABLE_NAME_COLUMN_KEY = "TABNAME";
34+
35+
public Db2JDBCCaseResolver(String sourceType)
36+
{
37+
super(sourceType, FederationSDKCasingMode.UPPER, FederationSDKCasingMode.LOWER);
38+
}
39+
40+
@Override
41+
protected String getCaseInsensitivelySchemaNameQueryTemplate()
42+
{
43+
return SCHEMA_NAME_QUERY_TEMPLATE;
44+
}
45+
46+
@Override
47+
protected String getCaseInsensitivelySchemaNameColumnKey()
48+
{
49+
return SCHEMA_NAME_COLUMN_KEY;
50+
}
51+
52+
@Override
53+
protected List<String> getCaseInsensitivelyTableNameQueryTemplate()
54+
{
55+
return List.of(TABLE_NAME_QUERY_TEMPLATE);
56+
}
57+
58+
@Override
59+
protected String getCaseInsensitivelyTableNameColumnKey()
60+
{
61+
return TABLE_NAME_COLUMN_KEY;
62+
}
63+
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import com.amazonaws.athena.connector.lambda.metadata.ListTablesRequest;
3838
import com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse;
3939
import com.amazonaws.athena.connector.lambda.security.FederatedIdentity;
40+
import com.amazonaws.athena.connectors.db2.resolver.Db2JDBCCaseResolver;
4041
import com.amazonaws.athena.connectors.jdbc.TestBase;
4142
import com.amazonaws.athena.connectors.jdbc.connection.DatabaseConnectionConfig;
4243
import com.amazonaws.athena.connectors.jdbc.connection.JdbcConnectionFactory;
@@ -96,7 +97,7 @@ public void setup() throws Exception {
9697
this.secretsManager = Mockito.mock(SecretsManagerClient.class);
9798
this.athena = Mockito.mock(AthenaClient.class);
9899
Mockito.when(this.secretsManager.getSecretValue(Mockito.eq(GetSecretValueRequest.builder().secretId("testSecret").build()))).thenReturn(GetSecretValueResponse.builder().secretString("{\"user\": \"testUser\", \"password\": \"testPassword\"}").build());
99-
this.db2MetadataHandler = new Db2MetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, this.jdbcConnectionFactory, com.google.common.collect.ImmutableMap.of());
100+
this.db2MetadataHandler = new Db2MetadataHandler(databaseConnectionConfig, this.secretsManager, this.athena, this.jdbcConnectionFactory, com.google.common.collect.ImmutableMap.of(), new Db2JDBCCaseResolver(Db2Constants.NAME));
100101
this.federatedIdentity = Mockito.mock(FederatedIdentity.class);
101102
this.blockAllocator = new BlockAllocatorImpl();
102103
}

‎athena-db2/src/test/java/com/amazonaws/athena/connectors/db2/resolver/Db2JDBCCaseResolverTest.java‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
/*-
2+
* #%L
3+
* athena-db2
4+
* %%
5+
* Copyright (C) 2019 - 2025 Amazon Web Services
6+
* %%
7+
* Licensed under the Apache License, Version 2.0 (the "License");
8+
* you may not use this file except in compliance with the License.
9+
* You may obtain a copy of the License at
10+
*
11+
* http://www.apache.org/licenses/LICENSE-2.0
12+
*
13+
* Unless required by applicable law or agreed to in writing, software
14+
* distributed under the License is distributed on an "AS IS" BASIS,
15+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16+
* See the License for the specific language governing permissions and
17+
* limitations under the License.
18+
* #L%
19+
*/
20+
package com.amazonaws.athena.connectors.db2.resolver;
21+
22+
import com.amazonaws.athena.connector.lambda.domain.TableName;
23+
import com.amazonaws.athena.connector.lambda.resolver.CaseResolver;
24+
import com.amazonaws.athena.connectors.db2.Db2Constants;
25+
import com.amazonaws.athena.connectors.jdbc.TestBase;
26+
import com.amazonaws.athena.connectors.jdbc.resolver.DefaultJDBCCaseResolver;
27+
import org.junit.Before;
28+
import org.junit.Test;
29+
import org.mockito.Mockito;
30+
31+
import java.sql.Connection;
32+
import java.sql.PreparedStatement;
33+
import java.sql.ResultSet;
34+
import java.sql.SQLException;
35+
import java.sql.Types;
36+
import java.util.Map;
37+
import java.util.concurrent.atomic.AtomicInteger;
38+
39+
import static com.amazonaws.athena.connector.lambda.resolver.CaseResolver.CASING_MODE_CONFIGURATION_KEY;
40+
import static org.junit.Assert.assertEquals;
41+
import static org.mockito.ArgumentMatchers.any;
42+
import static org.mockito.Mockito.when;
43+
44+
public class Db2JDBCCaseResolverTest extends TestBase
45+
{
46+
private Connection mockConnection;
47+
private PreparedStatement preparedStatement;
48+
49+
@Before
50+
public void setup() throws SQLException
51+
{
52+
mockConnection = Mockito.mock(Connection.class);
53+
preparedStatement = Mockito.mock(PreparedStatement.class);
54+
when(mockConnection.prepareStatement(any())).thenReturn(preparedStatement);
55+
}
56+
57+
@Test
58+
public void testCaseInsensitiveCaseOnName() throws SQLException
59+
{
60+
String schemaName = "oRaNgE";
61+
String tableName = "ApPlE";
62+
DefaultJDBCCaseResolver resolver = new Db2JDBCCaseResolver(Db2Constants.NAME);
63+
64+
// Mock schema name result
65+
String[] schemaCols = {"SCHEMANAME"};
66+
int[] schemaTypes = {Types.VARCHAR};
67+
Object[][] schemaData = {{schemaName.toLowerCase()}};
68+
69+
ResultSet schemaResultSet = mockResultSet(schemaCols, schemaTypes, schemaData, new AtomicInteger(-1));
70+
when(preparedStatement.executeQuery()).thenReturn(schemaResultSet);
71+
72+
String adjustedSchemaName = resolver.getAdjustedSchemaNameString(mockConnection, schemaName, Map.of(
73+
CASING_MODE_CONFIGURATION_KEY, CaseResolver.FederationSDKCasingMode.CASE_INSENSITIVE_SEARCH.name()));
74+
assertEquals(schemaName.toLowerCase(), adjustedSchemaName);
75+
76+
// Mock table name result
77+
String[] tableCols = {"TABNAME"};
78+
int[] tableTypes = {Types.VARCHAR};
79+
Object[][] tableData = {{tableName.toUpperCase()}};
80+
81+
ResultSet tableResultSet = mockResultSet(tableCols, tableTypes, tableData, new AtomicInteger(-1));
82+
when(preparedStatement.executeQuery()).thenReturn(tableResultSet);
83+
84+
String adjustedTableName = resolver.getAdjustedTableNameString(mockConnection, schemaName, tableName, Map.of(
85+
CASING_MODE_CONFIGURATION_KEY, CaseResolver.FederationSDKCasingMode.CASE_INSENSITIVE_SEARCH.name()));
86+
assertEquals(tableName.toUpperCase(), adjustedTableName);
87+
}
88+
89+
@Test
90+
public void testCaseInsensitiveCaseOnObject() throws SQLException
91+
{
92+
String schemaName = "oRaNgE";
93+
String tableName = "ApPlE";
94+
DefaultJDBCCaseResolver resolver = new Db2JDBCCaseResolver(Db2Constants.NAME);
95+
96+
// Mock schema and table result sets
97+
ResultSet schemaResultSet = mockResultSet(
98+
new String[]{"SCHEMANAME"},
99+
new int[]{Types.VARCHAR},
100+
new Object[][]{{schemaName.toLowerCase()}},
101+
new AtomicInteger(-1));
102+
103+
ResultSet tableResultSet = mockResultSet(
104+
new String[]{"TABNAME"},
105+
new int[]{Types.VARCHAR},
106+
new Object[][]{{tableName.toUpperCase()}},
107+
new AtomicInteger(-1));
108+
109+
when(preparedStatement.executeQuery()).thenReturn(schemaResultSet).thenReturn(tableResultSet);
110+
111+
TableName adjusted = resolver.getAdjustedTableNameObject(
112+
mockConnection,
113+
new TableName(schemaName, tableName),
114+
Map.of(CASING_MODE_CONFIGURATION_KEY, CaseResolver.FederationSDKCasingMode.CASE_INSENSITIVE_SEARCH.name()));
115+
assertEquals(new TableName(schemaName.toLowerCase(), tableName.toUpperCase()), adjusted);
116+
}
117+
}

0 commit comments

Comments
 (0)