Skip to content

Commit

Permalink
Use stringbuilder to construct default jdbc connection string and add…
Browse files Browse the repository at this point in the history
… unit tests
  • Loading branch information
ejeffrli committed Jan 10, 2025
1 parent 09e8689 commit 6b0636c
Show file tree
Hide file tree
Showing 6 changed files with 231 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -92,7 +92,7 @@ public Connection getGlueConnection(String glueConnectionName) throws RuntimeExc
}
}

private Map<String, String> authenticationConfigurationToMap(AuthenticationConfiguration auth)
public Map<String, String> authenticationConfigurationToMap(AuthenticationConfiguration auth)
{
Map<String, String> authMap = new HashMap<>();

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,59 @@
package com.amazonaws.athena.connector.lambda.connection;

/*-
* #%L
* Amazon Athena Query Federation SDK
* %%
* Copyright (C) 2019 Amazon Web Services
* %%
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
* #L%
*/

import software.amazon.awssdk.services.glue.model.AuthenticationConfiguration;

import org.junit.Test;

import java.util.HashMap;
import java.util.Map;

import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.KMS_KEY_ID;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.SECRET_NAME;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.SPILL_KMS_KEY_ID;
import static org.junit.Assert.*;

public class EnvironmentPropertiesTest
{
private static final String EXAMPLE_SECRET_ARN = "arn:aws:secretsmanager:us-east-1:012345678912:secret:secretname-CMyiKm";
private static final String EXAMPLE_KMS_KEY_ID = "1234abcd-12ab-34cd-56ef-1234567890ab";

@Test
public void authenticationConfigurationToMapTest()
{
AuthenticationConfiguration auth = AuthenticationConfiguration.builder().secretArn(EXAMPLE_SECRET_ARN).build();
Map<String, String> authMap = new EnvironmentProperties().authenticationConfigurationToMap(auth);
assertEquals("secretname", authMap.get(SECRET_NAME));
}

@Test
public void athenaPropertiesToEnvironmentTest()
{
Map<String, String> athenaProperties = new HashMap<>();
athenaProperties.put(SPILL_KMS_KEY_ID, EXAMPLE_KMS_KEY_ID);
athenaProperties = new EnvironmentProperties().athenaPropertiesToEnvironment(athenaProperties);

assertFalse(athenaProperties.containsKey(SPILL_KMS_KEY_ID));
assertTrue(athenaProperties.containsKey(KMS_KEY_ID));
assertEquals(EXAMPLE_KMS_KEY_ID, athenaProperties.get(KMS_KEY_ID));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -39,10 +39,11 @@ public Map<String, String> connectionPropertiesToEnvironment(Map<String, String>
HashMap<String, String> environment = new HashMap<>();

// now construct jdbc string
String connectionString = getConnectionStringPrefix(connectionProperties) + connectionProperties.get(HOST)
+ ":" + connectionProperties.get(PORT) + getDatabase(connectionProperties) + getJdbcParameters(connectionProperties);
StringBuilder connectionString = new StringBuilder(getConnectionStringPrefix(connectionProperties));
connectionString.append(connectionProperties.get(HOST)).append(":").append(connectionProperties.get(PORT));
connectionString.append(getDatabase(connectionProperties)).append(getJdbcParameters(connectionProperties));

environment.put(DEFAULT, connectionString);
environment.put(DEFAULT, connectionString.toString());
return environment;
}

Expand All @@ -55,16 +56,17 @@ protected String getDatabase(Map<String, String> connectionProperties)

protected String getJdbcParameters(Map<String, String> connectionProperties)
{
String params = getJdbcParametersSeparator() + connectionProperties.getOrDefault(JDBC_PARAMS, "");
StringBuilder params = new StringBuilder(getJdbcParametersSeparator());
params.append(connectionProperties.getOrDefault(JDBC_PARAMS, ""));

if (connectionProperties.containsKey(SECRET_NAME)) {
if (connectionProperties.containsKey(JDBC_PARAMS)) { // need to add delimiter
params = params + getDelimiter();
params.append(getDelimiter());
}
params = params + "${" + connectionProperties.get(SECRET_NAME) + "}";
params.append("${").append(connectionProperties.get(SECRET_NAME)).append("}");
}

return params;
return params.toString();
}

protected String getDatabaseSeparator()
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
package com.amazonaws.athena.connectors.mysql;

import org.junit.Before;
import org.junit.Test;

import java.util.HashMap;
import java.util.Map;

import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.DATABASE;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.DEFAULT;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.HOST;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.JDBC_PARAMS;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.PORT;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.SECRET_NAME;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

public class MySqlEnvironmentPropertiesTest
{
Map<String, String> connectionProperties;
MySqlEnvironmentProperties mySqlEnvironmentProperties;

@Before
public void setUp()
{
connectionProperties = new HashMap<>();
connectionProperties.put(HOST, "mysql.host");
connectionProperties.put(PORT, "3306");
connectionProperties.put(DATABASE, "test");
connectionProperties.put(JDBC_PARAMS, "key=value&key2=value2");
connectionProperties.put(SECRET_NAME, "secret");
mySqlEnvironmentProperties = new MySqlEnvironmentProperties();
}

@Test
public void connectionPropertiesWithJdbcParams()
{
String connectionString = "mysql://jdbc:mysql://mysql.host:3306/test?key=value&key2=value2&${secret}";
Map<String, String> mysqlConnectionProperties = mySqlEnvironmentProperties.connectionPropertiesToEnvironment(connectionProperties);
assertTrue(mysqlConnectionProperties.containsKey(DEFAULT));
assertEquals(connectionString, mysqlConnectionProperties.get(DEFAULT));
}

@Test
public void connectionPropertiesWithoutJdbcParams()
{
Map<String, String> noJdbcParams = new HashMap<>(connectionProperties);
noJdbcParams.remove(JDBC_PARAMS);
String connectionString = "mysql://jdbc:mysql://mysql.host:3306/test?${secret}";
Map<String, String> mysqlConnectionProperties = mySqlEnvironmentProperties.connectionPropertiesToEnvironment(noJdbcParams);
assertTrue(mysqlConnectionProperties.containsKey(DEFAULT));
assertEquals(connectionString, mysqlConnectionProperties.get(DEFAULT));
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
package com.amazonaws.athena.connectors.postgresql;

import org.junit.Before;
import org.junit.Test;

import java.util.HashMap;
import java.util.Map;

import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.DATABASE;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.DEFAULT;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.HOST;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.JDBC_PARAMS;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.PORT;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.SECRET_NAME;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

public class PostGreSqlEnvironmentPropertiesTest
{
Map<String, String> connectionProperties;
PostGreSqlEnvironmentProperties postGreSqlEnvironmentProperties;

@Before
public void setUp()
{
connectionProperties = new HashMap<>();
connectionProperties.put(HOST, "postgres.host");
connectionProperties.put(PORT, "5432");
connectionProperties.put(DATABASE, "default");
connectionProperties.put(JDBC_PARAMS, "key=value&key2=value2");
connectionProperties.put(SECRET_NAME, "secret");
postGreSqlEnvironmentProperties = new PostGreSqlEnvironmentProperties();
}

@Test
public void connectionPropertiesWithJdbcParams()
{
String connectionString = "postgres://jdbc:postgresql://postgres.host:5432/default?key=value&key2=value2&${secret}";
Map<String, String> pgsqlConnectionProperties = postGreSqlEnvironmentProperties.connectionPropertiesToEnvironment(connectionProperties);
assertTrue(pgsqlConnectionProperties.containsKey(DEFAULT));
assertEquals(connectionString, pgsqlConnectionProperties.get(DEFAULT));
}

@Test
public void connectionPropertiesWithoutJdbcParams()
{
Map<String, String> noJdbcParams = new HashMap<>(connectionProperties);
noJdbcParams.remove(JDBC_PARAMS);
String connectionString = "postgres://jdbc:postgresql://postgres.host:5432/default?${secret}";
Map<String, String> pgsqlConnectionProperties = postGreSqlEnvironmentProperties.connectionPropertiesToEnvironment(noJdbcParams);
assertTrue(pgsqlConnectionProperties.containsKey(DEFAULT));
assertEquals(connectionString, pgsqlConnectionProperties.get(DEFAULT));
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
package com.amazonaws.athena.connectors.redshift;

import org.junit.Before;
import org.junit.Test;

import java.util.HashMap;
import java.util.Map;

import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.DATABASE;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.DEFAULT;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.HOST;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.JDBC_PARAMS;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.PORT;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.SECRET_NAME;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertTrue;

public class RedshiftEnvironmentPropertiesTest
{
Map<String, String> connectionProperties;
RedshiftEnvironmentProperties redshiftEnvironmentProperties;

@Before
public void setUp()
{
connectionProperties = new HashMap<>();
connectionProperties.put(HOST, "redshift.host");
connectionProperties.put(PORT, "5439");
connectionProperties.put(DATABASE, "default");
connectionProperties.put(JDBC_PARAMS, "key=value&key2=value2");
connectionProperties.put(SECRET_NAME, "secret");
redshiftEnvironmentProperties = new RedshiftEnvironmentProperties();
}

@Test
public void connectionPropertiesWithJdbcParams()
{
String connectionString = "redshift://jdbc:redshift://redshift.host:5439/default?key=value&key2=value2&${secret}";
Map<String, String> mysqlConnectionProperties = redshiftEnvironmentProperties.connectionPropertiesToEnvironment(connectionProperties);
assertTrue(mysqlConnectionProperties.containsKey(DEFAULT));
assertEquals(connectionString, mysqlConnectionProperties.get(DEFAULT));
}

@Test
public void connectionPropertiesWithoutJdbcParams()
{
Map<String, String> noJdbcParams = new HashMap<>(connectionProperties);
noJdbcParams.remove(JDBC_PARAMS);
String connectionString = "redshift://jdbc:redshift://redshift.host:5439/default?${secret}";
Map<String, String> mysqlConnectionProperties = redshiftEnvironmentProperties.connectionPropertiesToEnvironment(noJdbcParams);
assertTrue(mysqlConnectionProperties.containsKey(DEFAULT));
assertEquals(connectionString, mysqlConnectionProperties.get(DEFAULT));
}
}

0 comments on commit 6b0636c

Please sign in to comment.