Feature/okta based redshift plugin

build.gradle

@@ -32,9 +32,11 @@ dependencies {
     }
     def withoutslf4jSupport = { exclude group: 'org.slf4j', module: 'slf4j-api' }
 
+    implementation 'com.amazonaws:aws-java-sdk-core:1.12.731', withoutjmespathSupport
     implementation 'com.amazonaws:aws-java-sdk-redshift:1.12.731', withoutjmespathSupport
     implementation 'com.amazonaws:aws-java-sdk-sts:1.12.731', withoutjmespathSupport
     implementation 'com.amazonaws:aws-java-sdk-redshiftserverless:1.12.731', withoutjmespathSupport
+    implementation 'com.amazonaws:aws-java-sdk-sso:1.12.731'
     implementation 'org.apache.httpcomponents:httpclient:4.5.14'
     implementation 'com.fasterxml.jackson.core:jackson-databind:2.16.0'
     implementation 'com.fasterxml.jackson.core:jackson-core:2.16.0'
@@ -73,13 +75,16 @@ processResources {
 }
 
 jar {
+    from {
+        configurations.runtimeClasspath.collect { it.isDirectory() ? it : zipTree(it) }
+    }
+    duplicatesStrategy = DuplicatesStrategy.EXCLUDE
     manifest {
         attributes("Automatic-Module-Name": 'com.amazon.redshift.jdbc')
         attributes("Main-Class": "com.amazon.redshift.util.RedshiftJDBCMain")
         attributes("Specification-Title": "JDBC")
         attributes("Specification-Version": "4.2")
         attributes("Specification-Vendor": "Oracle Corporation")
-        attributes("Class-Path": configurations.runtimeClasspath.collect { it.getName() }.join(' '))
     }
 }
 

src/main/java/com/amazon/redshift/Driver.java

@@ -480,7 +480,11 @@ public Connection getResult(long timeout) throws SQLException {
    * @throws SQLException if the connection could not be made
    */
   private static Connection makeConnection(String url, RedshiftProperties props, RedshiftLogger logger) throws SQLException {
-    return new RedshiftConnectionImpl(hostSpecs(props), user(props), database(props), props, url, logger);
+
+      String iamauth = props.getProperty("iamauth");
+      System.out.println(iamauth);
+
+      return new RedshiftConnectionImpl(hostSpecs(props), user(props), database(props), props, url, logger);
   }
 
   /**
@@ -645,7 +649,10 @@ public static RedshiftProperties parseURL(String url, RedshiftProperties default
       	urlArgs = queryString;
     } // IAM
     else {
-    	urlProps.setProperty(RedshiftProperty.IAM_AUTH.getName(), String.valueOf(iamAuth));
+    	// Only set iamAuth to false if it's not already explicitly set by the user
+    	if (urlProps.getProperty(RedshiftProperty.IAM_AUTH.getName()) == null) {
+    		urlProps.setProperty(RedshiftProperty.IAM_AUTH.getName(), String.valueOf(iamAuth));
+    	}
 
 	    if (urlServer.startsWith("//")) {
 	      urlServer = urlServer.substring(2);

src/main/java/com/amazon/redshift/TestOktaDriver.java

@@ -0,0 +1,47 @@
+package com.amazon.redshift;
+
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.ResultSet;
+import java.sql.Statement;
+import java.util.Properties;
+
+public class TestOktaDriver {
+    public static void main(String[] args) throws Exception {
+        // Register the driver
+        Class.forName("com.amazon.redshift.Driver");
+
+        // Connection URL
+        String url = "jdbc:redshift://hubble.chg6aanrjt24.eu-west-1.redshift.amazonaws.com:5439/dev";
+
+        // Connection properties
+        Properties props = new Properties();
+        props.setProperty("plugin_name", "com.amazon.redshift.plugin.OktaRedshiftPlugin");
+        props.setProperty("iamauth", "true");
+
+        // Plugin parameters
+        props.setProperty("ssoRoleName", "OktaDataViewer");
+        props.setProperty("region", "eu-north-1");
+        props.setProperty("ssoStartUrl", "https://d-c3672deb5f.awsapps.com/start");
+        props.setProperty("preferred_role", "hubble-rbac/DataViewer");
+        props.setProperty("ssoAccountID", "899945594626");
+        props.setProperty("clusterid", "hubble");
+
+        // Test connection
+        try (Connection conn = DriverManager.getConnection(url, props)) {
+            System.out.println("Connected successfully!");
+
+            // Test query
+            try (Statement stmt = conn.createStatement()) {
+                ResultSet rs = stmt.executeQuery("SELECT current_user, current_database()");
+                if (rs.next()) {
+                    System.out.println("Current user: " + rs.getString(1));
+                    System.out.println("Current database: " + rs.getString(2));
+                }
+            }
+        } catch (Exception e) {
+            System.err.println("Connection failed: " + e.getMessage());
+            e.printStackTrace();
+        }
+    }
+}

src/main/java/com/amazon/redshift/core/IamHelper.java

@@ -36,6 +36,7 @@
 import com.amazon.redshift.logger.LogLevel;
 import com.amazon.redshift.logger.RedshiftLogger;
 import com.amazon.redshift.plugin.utils.RequestUtils;
+import com.amazon.redshift.plugin.OktaRedshiftPlugin;
 import com.amazon.redshift.util.GT;
 import com.amazon.redshift.util.RedshiftException;
 import com.amazon.redshift.util.RedshiftState;
@@ -65,6 +66,7 @@ public final class IamHelper extends IdpAuthHelper {
   public static final int GET_CLUSTER_CREDENTIALS_SAML_V2_API = 3;
   public static final int GET_CLUSTER_CREDENTIALS_JWT_V2_API = 4;
   public static final int GET_SERVERLESS_CREDENTIALS_V1_API = 5;
+  public static final int GET_CLUSTER_CREDENTIALS_PLUGIN_DIRECT = 6;
 
   private static final Pattern HOST_PATTERN =
           Pattern.compile("(.+)\\.(.+)\\.(.+).redshift(-dev)?\\.amazonaws\\.com(.)*");
@@ -598,6 +600,7 @@ else if (RedshiftProperty.DB_GROUPS_FILTER.getName().equalsIgnoreCase(pluginArgK
 
         settings.m_idpToken = idpToken;
       }
+
     } // Group federation API for plugin
 
     setClusterCredentials(provider, settings, log, providerType, idpCredentialsRefresh, getClusterCredentialApiType);
@@ -682,6 +685,29 @@ private static void setClusterCredentials(AWSCredentialsProvider credProvider, R
             log.logInfo(now + ": Using GetClusterCredentialsResultV2 with TimeToRefresh " + iamResult.getNextRefreshTime());
           }
 
+          break;
+
+        case GET_CLUSTER_CREDENTIALS_PLUGIN_DIRECT:
+          // Plugin directly provides database credentials
+          if (RedshiftLogger.isEnable())
+            log.log(LogLevel.DEBUG, "Using plugin-provided database credentials directly");
+
+          // Get the credentials from the provider
+          AWSCredentials pluginCredentials = credProvider.getCredentials();
+          if (pluginCredentials instanceof OktaRedshiftPlugin.DatabaseCredentials) {
+            OktaRedshiftPlugin.DatabaseCredentials dbCreds = (OktaRedshiftPlugin.DatabaseCredentials) pluginCredentials;
+            settings.m_username = dbCreds.getAWSAccessKeyId(); // username stored in access key field
+            settings.m_password = dbCreds.getAWSSecretKey(); // password stored in secret key field
+
+            if (RedshiftLogger.isEnable()) {
+              Date now = new Date();
+              log.logInfo(now + ": Using plugin database credentials with expiration " + dbCreds.getExpiration());
+            }
+          } else {
+            throw new RedshiftException("Expected DatabaseCredentials from plugin but got: " + 
+                pluginCredentials.getClass().getSimpleName(), RedshiftState.UNEXPECTED_ERROR);
+          }
+
           break;
       }
     }
@@ -1212,6 +1238,13 @@ static String getCredentialsV2CacheKey(RedshiftJDBCSettings settings, Credential
   private static int findTypeOfGetClusterCredentialsAPI(RedshiftJDBCSettings settings,
       CredentialProviderType providerType, AWSCredentialsProvider provider) {
 
+    // set the GET_CLUSTER_CREDENTIALS_PLUGIN_DIRECT for OktaRedshiftPlugin that returns DatabaseCredentials directly
+    if (providerType == CredentialProviderType.PLUGIN && 
+        settings.m_credentialsProvider != null &&
+        settings.m_credentialsProvider.contains("OktaRedshiftPlugin")) {
+      return GET_CLUSTER_CREDENTIALS_PLUGIN_DIRECT;
+    }
+
     if (!settings.m_isServerless)
     {
       if (!settings.m_groupFederation)

src/main/java/com/amazon/redshift/plugin/CommonCredentialsProvider.java

@@ -22,6 +22,7 @@
 import org.apache.commons.logging.LogFactory;
 
 import java.io.IOException;
+import java.net.URISyntaxException;
 import java.net.URL;
 import java.util.Collections;
 import java.util.Enumeration;
@@ -129,7 +130,7 @@ public NativeTokenHolder getCredentials() throws RedshiftException {
         return credentials;
     }
 
-    protected abstract NativeTokenHolder getAuthToken() throws IOException;
+    protected abstract NativeTokenHolder getAuthToken() throws IOException, URISyntaxException;
 
     @Override
     public void refresh() throws RedshiftException {