From cc02fa4b5394596dd8f6490261f96fdf83908c07 Mon Sep 17 00:00:00 2001
From: Andrey Kuzin <akuzin@amazon.com>
Date: Tue, 16 Jan 2024 23:47:08 +0000
Subject: [PATCH] Bug Fixes Glue Integration

---
 .../athena/connector/lambda/GlueConnectionUtils.java         | 5 ++++-
 .../jdbc/connection/DatabaseConnectionConfigBuilder.java     | 2 +-
 .../amazonaws/athena/connectors/jdbc/manager/JDBCUtil.java   | 4 +---
 3 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/GlueConnectionUtils.java b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/GlueConnectionUtils.java
index 8697981b1c..c2ff2c5990 100644
--- a/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/GlueConnectionUtils.java
+++ b/athena-federation-sdk/src/main/java/com/amazonaws/athena/connector/lambda/GlueConnectionUtils.java
@@ -51,9 +51,10 @@ public static Map<String, String> getGlueConnection()
         String glueConnectionName = envConfig.get(DEFAULT_GLUE_CONNECTION);
         if (glueConnectionName != null) {
             HashMap<String, String> cachedConfig = connectionNameCache.get(glueConnectionName);
-            if (cachedConfig != null) {
+            if (cachedConfig == null) {
                 try {
                     HashMap<String, HashMap<String, String>> athenaDriverPropertiesToMap = new HashMap<String, HashMap<String, String>>();
+
                     AWSGlue awsGlue = AWSGlueClientBuilder.standard().withClientConfiguration(new ClientConfiguration().withConnectionTimeout(CONNECT_TIMEOUT)).build();
                     GetConnectionResult glueConnection = awsGlue.getConnection(new GetConnectionRequest().withName(glueConnectionName));
                     Connection connection = glueConnection.getConnection();
@@ -74,6 +75,8 @@ public static Map<String, String> getGlueConnection()
                 catch (Exception err) {
                     logger.error("Error thrown during fetching of {} connection properties. {}", glueConnectionName, err.toString());
                 }
+            }else{
+                return cachedConfig;
             }
         }
         else {            
diff --git a/athena-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/connection/DatabaseConnectionConfigBuilder.java b/athena-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/connection/DatabaseConnectionConfigBuilder.java
index b2ebb576a8..e3e53c5f77 100644
--- a/athena-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/connection/DatabaseConnectionConfigBuilder.java
+++ b/athena-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/connection/DatabaseConnectionConfigBuilder.java
@@ -96,7 +96,7 @@ public List<DatabaseConnectionConfig> build()
         List<DatabaseConnectionConfig> databaseConnectionConfigs = new ArrayList<>();
 
         int numberOfCatalogs = 0;
-        if (StringUtils.isBlank(properties.get(DEFAULT_GLUE_CONNECTION))) {
+        if (!StringUtils.isBlank(properties.get(DEFAULT_GLUE_CONNECTION))) {
             databaseConnectionConfigs.add(extractDatabaseGlueConnectionConfig(DEFAULT_CONNECTION_STRING_PROPERTY));
             numberOfCatalogs++;
         }
diff --git a/athena-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/manager/JDBCUtil.java b/athena-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/manager/JDBCUtil.java
index dbd9f4d421..e5e59d3bac 100644
--- a/athena-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/manager/JDBCUtil.java
+++ b/athena-jdbc/src/main/java/com/amazonaws/athena/connectors/jdbc/manager/JDBCUtil.java
@@ -54,12 +54,10 @@ public static DatabaseConnectionConfig getSingleDatabaseConfigFromEnv(String dat
         List<DatabaseConnectionConfig> databaseConnectionConfigs = DatabaseConnectionConfigBuilder.buildFromSystemEnv(databaseEngine, configOptions);
 
         for (DatabaseConnectionConfig databaseConnectionConfig : databaseConnectionConfigs) {
-            if (DatabaseConnectionConfigBuilder.DEFAULT_CONNECTION_STRING_PROPERTY.equals(databaseConnectionConfig.getCatalog())
-                    && databaseEngine.equals(databaseConnectionConfig.getEngine())) {
+            if (DatabaseConnectionConfigBuilder.DEFAULT_CONNECTION_STRING_PROPERTY.equals(databaseConnectionConfig.getCatalog())) {
                 return databaseConnectionConfig;
             }
         }
-
         throw new RuntimeException(String.format("Must provide default connection string parameter %s for database type %s",
                 DatabaseConnectionConfigBuilder.DEFAULT_CONNECTION_STRING_PROPERTY, databaseEngine));
     }