Skip to content

Commit

Permalink
Fix lambda issues with glue (awslabs#2308)
Browse files Browse the repository at this point in the history
  • Loading branch information
@ejeffrli
ejeffrli authored Oct 2, 2024
1 parent 63a6075 commit a787426
Show file tree
Hide file tree
Showing 4 changed files with 54 additions and 75 deletions.
8 changes: 6 additions & 2 deletions ...src/main/java/com/amazonaws/athena/connector/lambda/connection/EnvironmentProperties.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import software.amazon.awssdk.services.glue.model.GetConnectionResponse;

import java.time.Duration;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

Expand All @@ -51,7 +52,7 @@ public Map<String, String> createEnvironment() throws RuntimeException
HashMap<String, String> connectionEnvironment = new HashMap<>();
if (StringUtils.isNotBlank(glueConnectionName)) {
Connection connection = getGlueConnection(glueConnectionName);
Map<String, String> connectionProperties = connection.connectionPropertiesAsStrings();
Map<String, String> connectionProperties = new HashMap<>(connection.connectionPropertiesAsStrings());
connectionProperties.putAll(authenticationConfigurationToMap(connection.authenticationConfiguration()));

connectionEnvironment.putAll(connectionPropertiesToEnvironment(connectionProperties));
Expand Down Expand Up @@ -86,7 +87,10 @@ private Map<String, String> authenticationConfigurationToMap(AuthenticationConfi

if (StringUtils.isNotBlank(auth.secretArn())) {
String[] splitArn = auth.secretArn().split(":");
authMap.put(SECRET_NAME, splitArn[splitArn.length - 1]);
String[] secretNameWithRandom = splitArn[splitArn.length - 1].split("-"); // 6 random characters at end. at least length of 2
String[] secretNameArray = Arrays.copyOfRange(secretNameWithRandom, 0, secretNameWithRandom.length - 1);
String secretName = String.join("-", secretNameArray); // add back the dashes
authMap.put(SECRET_NAME, secretName);
}
return authMap;
}
Expand Down
60 changes: 22 additions & 38 deletions ...java/com/amazonaws/athena/connectors/jdbc/connection/DatabaseConnectionConfigBuilder.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,10 +43,6 @@ public class DatabaseConnectionConfigBuilder
private static final String SECRET_PATTERN_STRING = "\\$\\{(([a-z-]+!)?[a-zA-Z0-9:/_+=.@-]+)}";
public static final Pattern SECRET_PATTERN = Pattern.compile(SECRET_PATTERN_STRING);

// Config variables used when glue connection supplements connection properties
public static final String DEFAULT_JDBC_CONNECTION_URL_PROPERTY = "default_connection_string";
public static final String DEFAULT_SECRET_PROPERTY = "secret_name";

public static final String DEFAULT_GLUE_CONNECTION = "glue_connection";

private Map<String, String> properties;
Expand Down Expand Up @@ -93,37 +89,33 @@ public DatabaseConnectionConfigBuilder properties(final Map<String, String> prop
public List<DatabaseConnectionConfig> build()
{
Validate.notEmpty(this.properties, "properties must not be empty");
Validate.isTrue(properties.containsKey(DEFAULT_CONNECTION_STRING_PROPERTY) || properties.containsKey(DEFAULT_JDBC_CONNECTION_URL_PROPERTY), "Default connection string must be present");
Validate.isTrue(properties.containsKey(DEFAULT_CONNECTION_STRING_PROPERTY), "Default connection string must be present");

List<DatabaseConnectionConfig> databaseConnectionConfigs = new ArrayList<>();

int numberOfCatalogs = 0;
if (!StringUtils.isBlank(properties.get(DEFAULT_GLUE_CONNECTION))) {
databaseConnectionConfigs.add(extractDatabaseGlueConnectionConfig(DEFAULT_CONNECTION_STRING_PROPERTY));
numberOfCatalogs++;
}
else {
for (Map.Entry<String, String> property : this.properties.entrySet()) {
final String key = property.getKey();
final String value = property.getValue();

String catalogName;
if (DEFAULT_CONNECTION_STRING_PROPERTY.equals(key.toLowerCase())) {
catalogName = key.toLowerCase();
}
else if (key.endsWith(CONNECTION_STRING_PROPERTY_SUFFIX)) {
catalogName = key.replace(CONNECTION_STRING_PROPERTY_SUFFIX, "");
}
else {
// unknown property ignore
continue;
}
databaseConnectionConfigs.add(extractDatabaseConnectionConfig(catalogName, value));
for (Map.Entry<String, String> property : this.properties.entrySet()) {
final String key = property.getKey();
final String value = property.getValue();

numberOfCatalogs++;
if (numberOfCatalogs > MUX_CATALOG_LIMIT) {
throw new RuntimeException("Too many database instances in mux. Max supported is " + MUX_CATALOG_LIMIT);
}
String catalogName;
if (DEFAULT_CONNECTION_STRING_PROPERTY.equals(key.toLowerCase())) {
catalogName = key.toLowerCase();
}
else if (key.endsWith(CONNECTION_STRING_PROPERTY_SUFFIX)) {
catalogName = key.replace(CONNECTION_STRING_PROPERTY_SUFFIX, "");
}
else {
// unknown property ignore
continue;
}
databaseConnectionConfigs.add(extractDatabaseConnectionConfig(catalogName, value));

if (StringUtils.isBlank(properties.get(DEFAULT_GLUE_CONNECTION))) {
numberOfCatalogs++; // Mux is not supported with glue. Do not count
}
if (numberOfCatalogs > MUX_CATALOG_LIMIT) {
throw new RuntimeException("Too many database instances in mux. Max supported is " + MUX_CATALOG_LIMIT);
}
}

Expand Down Expand Up @@ -153,14 +145,6 @@ private DatabaseConnectionConfig extractDatabaseConnectionConfig(final String ca
.orElseGet(() -> new DatabaseConnectionConfig(catalogName, this.engine, jdbcConnectionString));
}

private DatabaseConnectionConfig extractDatabaseGlueConnectionConfig(final String catalogName)
{
final String jdbcConnectionString = properties.get(DEFAULT_JDBC_CONNECTION_URL_PROPERTY);
final String secretName = properties.get(DEFAULT_SECRET_PROPERTY);
Validate.notBlank(jdbcConnectionString, "JDBC Connection string must not be blank.");
return StringUtils.isBlank(secretName) ? new DatabaseConnectionConfig(catalogName, this.engine, jdbcConnectionString) : new DatabaseConnectionConfig(catalogName, this.engine, jdbcConnectionString, secretName);
}

private Optional<String> extractSecretName(final String jdbcConnectionString)
{
Matcher secretMatcher = SECRET_PATTERN.matcher(jdbcConnectionString);
Expand Down
35 changes: 0 additions & 35 deletions .../com/amazonaws/athena/connectors/jdbc/connection/DatabaseConnectionConfigBuilderTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -125,40 +125,5 @@ public void validSecretsSyntaxTest()
Assert.assertEquals(secrets[i], databaseConnectionConfigs.get(i).getSecret());
}
}

@Test
public void buildUsingGlueConnectionWithSecret()
{
DatabaseConnectionConfig glueSupplementedConnection = new DatabaseConnectionConfig("default", "postgres",
"jdbc:postgresql://hostname/test", "testSecret");

List<DatabaseConnectionConfig> databaseConnectionConfigs = new DatabaseConnectionConfigBuilder()
.engine("postgres")
.properties(ImmutableMap.of(
"default", CONNECTION_STRING2,
"default_connection_string", CONNECTION_STRING5,
"secret_name", CONNECTION_STRING5_SECRET,
"glue_connection", MOCK_GLUE_CONNECTION_NAME))
.build();

Assert.assertEquals(Arrays.asList(glueSupplementedConnection), databaseConnectionConfigs);
}

@Test
public void buildUsingGlueConnectionNoSecret()
{
DatabaseConnectionConfig glueSupplementedConnection = new DatabaseConnectionConfig("default", "postgres",
"jdbc:postgresql://hostname/test");

List<DatabaseConnectionConfig> databaseConnectionConfigs = new DatabaseConnectionConfigBuilder()
.engine("postgres")
.properties(ImmutableMap.of(
"default", CONNECTION_STRING2,
"default_connection_string", CONNECTION_STRING5,
"glue_connection", MOCK_GLUE_CONNECTION_NAME))
.build();

Assert.assertEquals(Arrays.asList(glueSupplementedConnection), databaseConnectionConfigs);
}
}

26 changes: 26 additions & 0 deletions ...c/main/java/com/amazonaws/athena/connectors/snowflake/SnowflakeEnvironmentProperties.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,14 +21,34 @@

import com.amazonaws.athena.connectors.jdbc.JdbcEnvironmentProperties;

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.PORT;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.SCHEMA;
import static com.amazonaws.athena.connector.lambda.connection.EnvironmentConstants.WAREHOUSE;

public class SnowflakeEnvironmentProperties extends JdbcEnvironmentProperties
{
@Override
public Map<String, String> connectionPropertiesToEnvironment(Map<String, String> connectionProperties)
{
HashMap<String, String> environment = new HashMap<>();

// now construct jdbc string
String connectionString = getConnectionStringPrefix(connectionProperties) + connectionProperties.get(HOST);
if (connectionProperties.containsKey(PORT)) {
connectionString = connectionString + ":" + connectionProperties.get(PORT);
}
connectionString = connectionString + getDatabase(connectionProperties) + getJdbcParameters(connectionProperties);

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

@Override
protected String getConnectionStringPrefix(Map<String, String> connectionProperties)
{
Expand All @@ -47,4 +67,10 @@ protected String getDatabase(Map<String, String> connectionProperties)
+ "&schema=" + connectionProperties.get(SCHEMA);
return databaseString;
}

@Override
protected String getJdbcParametersSeparator()
{
return "&";
}
}

0 comments on commit a787426

Please sign in to comment.