[b/387249832] Move Redshift URL utilities to a Util class (#681)

misolt · web-flow · commit 2b3434600cfb · 2025-01-07T10:40:11.000+01:00
diff --git a/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/redshift/AbstractRedshiftConnector.java b/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/redshift/AbstractRedshiftConnector.java
@@ -17,9 +17,7 @@
 package com.google.edwmigration.dumper.application.dumper.connector.redshift;
 
 import com.google.common.base.Joiner;
-import com.google.common.collect.Iterables;
 import com.google.edwmigration.dumper.application.dumper.ConnectorArguments;
-import com.google.edwmigration.dumper.application.dumper.MetadataDumperUsageException;
 import com.google.edwmigration.dumper.application.dumper.annotations.RespectsArgumentDatabaseForConnection;
 import com.google.edwmigration.dumper.application.dumper.annotations.RespectsArgumentDriver;
 import com.google.edwmigration.dumper.application.dumper.annotations.RespectsArgumentHostUnlessUrl;
@@ -29,8 +27,8 @@
 import com.google.edwmigration.dumper.application.dumper.connector.AbstractJdbcConnector;
 import com.google.edwmigration.dumper.application.dumper.handle.Handle;
 import com.google.edwmigration.dumper.application.dumper.handle.JdbcHandle;
-import java.io.UnsupportedEncodingException;
 import java.sql.Driver;
+import java.sql.SQLException;
 import java.time.ZoneOffset;
 import java.time.format.DateTimeFormatter;
 import java.util.List;
@@ -69,12 +67,10 @@
 @RespectsArgumentJDBCUri
 public abstract class AbstractRedshiftConnector extends AbstractJdbcConnector {
 
-  @SuppressWarnings("UnusedVariable")
   private static final Logger LOG = LoggerFactory.getLogger(AbstractRedshiftConnector.class);
 
   protected static final DateTimeFormatter SQL_FORMAT =
       DateTimeFormatter.ISO_OFFSET_DATE_TIME.withZone(ZoneOffset.UTC);
-  public static final int OPT_PORT_DEFAULT = 5439;
 
   @Nonnull
   protected static CharSequence newWhereClause(List<String> clauseList, String... clauseArray) {
@@ -130,85 +126,15 @@ protected static CharSequence newWhereClause(List<String> clauseList, String...
    *
    * https://s3.amazonaws.com/redshift-downloads/drivers/jdbc/1.2.43.1067/Amazon+Redshift+JDBC+Driver+Install+Guide.pdf
    */
-  // this SHOuLD LAND UP IN THE ARGUMENT CLASS ...
-  @Nonnull
-  private static String requireNonNull(String val, String msg) throws MetadataDumperUsageException {
-    if (val != null) return val;
-    throw new MetadataDumperUsageException(msg);
-  }
-
-  @Nonnull
-  private String makeJdbcUrlPostgresql(ConnectorArguments arguments)
-      throws MetadataDumperUsageException, UnsupportedEncodingException {
-    String password = arguments.getPasswordIfFlagProvided().orElse(null);
-    return "jdbc:postgresql://"
-        + arguments.getHostOrDefault()
-        + ":"
-        + arguments.getPort(5439)
-        + "/"
-        + Iterables.getFirst(arguments.getDatabases(), "") //
-        + new JdbcPropBuilder("?=&")
-            .propOrWarn("user", arguments.getUser(), "--user must be specified")
-            .propOrWarn("password", password, "--password must be specified")
-            .prop("ssl", "true")
-            .toJdbcPart();
-  }
-
-  @Nonnull
-  private String makeJdbcUrlRedshiftSimple(ConnectorArguments arguments)
-      throws MetadataDumperUsageException, UnsupportedEncodingException {
-    String password = arguments.getPasswordIfFlagProvided().orElse(null);
-    return "jdbc:redshift://"
-        + arguments.getHostOrDefault()
-        + ":"
-        + arguments.getPort(5439)
-        + "/"
-        + Iterables.getFirst(arguments.getDatabases(), "") //
-        + new JdbcPropBuilder("?=&")
-            .propOrWarn("UID", arguments.getUser(), "--user must be specified")
-            .propOrWarn("PWD", password, "--password must be specified")
-            .toJdbcPart();
-  }
-
-  // TODO:  [cluster-id]:[region] syntax.
-  // either profile, or key+ secret
-  @Nonnull
-  private String makeJdbcUrlRedshiftIAM(ConnectorArguments arguments)
-      throws MetadataDumperUsageException, UnsupportedEncodingException {
-    String url =
-        "jdbc:redshift:iam://"
-            + arguments.getHostOrDefault()
-            + ":"
-            + arguments.getPort(5439)
-            + "/"
-            + Iterables.getFirst(arguments.getDatabases(), "");
-
-    if (arguments.getIAMProfile() != null)
-      url += new JdbcPropBuilder("?=&").prop("Profile", arguments.getIAMProfile()).toJdbcPart();
-    else if (arguments.getIAMAccessKeyID() != null && arguments.getIAMSecretAccessKey() != null)
-      url +=
-          new JdbcPropBuilder("?=&")
-              .prop("AccessKeyID", arguments.getIAMAccessKeyID())
-              .prop("SecretAccessKey", arguments.getIAMSecretAccessKey())
-              .propOrError("DbUser", arguments.getUser(), "--user must be specified")
-              .toJdbcPart();
-    // Will use the default IAM from ~/.aws/credentials/
-    // throw new MetadataDumperUsageException("Either --iam-profile or
-    // --iam-accesskeyid/--iam-secretaccesskey should be given");
-    return url;
-  }
 
   @Override
-  public Handle open(ConnectorArguments arguments) throws Exception {
+  @Nonnull
+  public Handle open(@Nonnull ConnectorArguments arguments) throws Exception {
 
     Driver driver =
         newDriver(
             arguments.getDriverPaths(), "com.amazon.redshift.jdbc.Driver", "org.postgresql.Driver");
 
-    //        LOG.debug("DRIVER IS " + driver.getClass().getCanonicalName());
-    //        LOG.debug("DRIVER CAN RS " + driver.acceptsURL("jdbc:redshift://host/db"));
-    //        LOG.debug("DRIVER CAN IAM " + driver.acceptsURL("jdbc:redshift:iam://host/db"));
-    //        LOG.debug("DRIVER CAN PG " + driver.acceptsURL("jdbc:postgresql://host/db"));
     String url = arguments.getUri();
     Optional<String> password = arguments.getPasswordIfFlagProvided();
     if (url == null) {
@@ -226,11 +152,11 @@ public Handle open(ConnectorArguments arguments) throws Exception {
               "The use of IAM authentication also requires the use of a Redshift-specific JDBC driver. Please use --"
                   + ConnectorArguments.OPT_DRIVER
                   + " to specify the path to the Redshift JDBC JAR, or use password authentication.");
-        url = makeJdbcUrlPostgresql(arguments);
+        url = RedshiftUrlUtil.makeJdbcUrlPostgresql(arguments);
       } else if (isAuthenticationPassword) {
-        url = makeJdbcUrlRedshiftSimple(arguments);
+        url = RedshiftUrlUtil.makeJdbcUrlRedshiftSimple(arguments);
       } else {
-        url = makeJdbcUrlRedshiftIAM(arguments);
+        url = RedshiftUrlUtil.makeJdbcUrlRedshiftIAM(arguments);
       }
     }
 
@@ -243,4 +169,11 @@ public Handle open(ConnectorArguments arguments) throws Exception {
 
     return JdbcHandle.newPooledJdbcHandle(dataSource, arguments.getThreadPoolSize());
   }
+
+  private static void logDriverInfo(@Nonnull Driver driver) throws SQLException {
+    LOG.debug("DRIVER IS " + driver.getClass().getCanonicalName());
+    LOG.debug("DRIVER CAN RS " + driver.acceptsURL("jdbc:redshift://host/db"));
+    LOG.debug("DRIVER CAN IAM " + driver.acceptsURL("jdbc:redshift:iam://host/db"));
+    LOG.debug("DRIVER CAN PG " + driver.acceptsURL("jdbc:postgresql://host/db"));
+  }
 }
diff --git a/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/redshift/RedshiftLogsConnector.java b/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/redshift/RedshiftLogsConnector.java
@@ -55,7 +55,7 @@
     arg = ConnectorArguments.OPT_PORT,
     description = "The port of the server.",
     required = ConnectorArguments.OPT_REQUIRED_IF_NOT_URL,
-    defaultValue = "" + RedshiftMetadataConnector.OPT_PORT_DEFAULT)
+    defaultValue = RedshiftUrlUtil.OPT_PORT_DEFAULT)
 @RespectsArgumentQueryLogDays
 @RespectsArgumentQueryLogStart
 @RespectsArgumentQueryLogEnd
diff --git a/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/redshift/RedshiftMetadataConnector.java b/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/redshift/RedshiftMetadataConnector.java
@@ -41,7 +41,7 @@
     arg = ConnectorArguments.OPT_PORT,
     description = "The port of the server.",
     required = ConnectorArguments.OPT_REQUIRED_IF_NOT_URL,
-    defaultValue = "" + RedshiftMetadataConnector.OPT_PORT_DEFAULT)
+    defaultValue = RedshiftUrlUtil.OPT_PORT_DEFAULT)
 public class RedshiftMetadataConnector extends AbstractRedshiftConnector
     implements MetadataConnector, RedshiftMetadataDumpFormat {
 
diff --git a/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/redshift/RedshiftRawLogsConnector.java b/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/redshift/RedshiftRawLogsConnector.java
@@ -63,7 +63,7 @@
     arg = ConnectorArguments.OPT_PORT,
     description = "The port of the server.",
     required = ConnectorArguments.OPT_REQUIRED_IF_NOT_URL,
-    defaultValue = "" + RedshiftMetadataConnector.OPT_PORT_DEFAULT)
+    defaultValue = RedshiftUrlUtil.OPT_PORT_DEFAULT)
 @RespectsArgumentAssessment
 @RespectsArgumentQueryLogDays
 @RespectsArgumentQueryLogStart
diff --git a/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/redshift/RedshiftUrlUtil.java b/dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/redshift/RedshiftUrlUtil.java
@@ -0,0 +1,103 @@
+/*
+ * Copyright 2022-2024 Google LLC
+ * Copyright 2013-2021 CompilerWorks
+ *
+ * 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.
+ */
+package com.google.edwmigration.dumper.application.dumper.connector.redshift;
+
+import com.google.common.collect.Iterables;
+import com.google.edwmigration.dumper.application.dumper.ConnectorArguments;
+import com.google.edwmigration.dumper.application.dumper.MetadataDumperUsageException;
+import java.io.UnsupportedEncodingException;
+import javax.annotation.Nonnull;
+import javax.annotation.ParametersAreNonnullByDefault;
+
+@ParametersAreNonnullByDefault
+class RedshiftUrlUtil {
+
+  private static final int DEFAULT_PORT = 5439;
+  // for annotations
+  static final String OPT_PORT_DEFAULT = "5439";
+
+  @Nonnull
+  static String makeJdbcUrlPostgresql(ConnectorArguments arguments)
+      throws MetadataDumperUsageException, UnsupportedEncodingException {
+    String password = arguments.getPasswordIfFlagProvided().orElse(null);
+    return makeScheme("postgresql")
+        + arguments.getHostOrDefault()
+        + ":"
+        + arguments.getPort(DEFAULT_PORT)
+        + "/"
+        + Iterables.getFirst(arguments.getDatabases(), "")
+        + new JdbcPropBuilder("?=&")
+            .propOrWarn("user", arguments.getUser(), "--user must be specified")
+            .propOrWarn("password", password, "--password must be specified")
+            .prop("ssl", "true")
+            .toJdbcPart();
+  }
+
+  @Nonnull
+  static String makeJdbcUrlRedshiftSimple(ConnectorArguments arguments)
+      throws MetadataDumperUsageException, UnsupportedEncodingException {
+    String password = arguments.getPasswordIfFlagProvided().orElse(null);
+    return makeScheme("redshift")
+        + arguments.getHostOrDefault()
+        + ":"
+        + arguments.getPort(DEFAULT_PORT)
+        + "/"
+        + Iterables.getFirst(arguments.getDatabases(), "")
+        + new JdbcPropBuilder("?=&")
+            .propOrWarn("UID", arguments.getUser(), "--user must be specified")
+            .propOrWarn("PWD", password, "--password must be specified")
+            .toJdbcPart();
+  }
+
+  // TODO:  [cluster-id]:[region] syntax.
+  // either profile, or key+ secret
+  @Nonnull
+  static String makeJdbcUrlRedshiftIAM(ConnectorArguments arguments)
+      throws UnsupportedEncodingException {
+    return makeScheme("redshift:iam")
+        + arguments.getHostOrDefault()
+        + ":"
+        + arguments.getPort(DEFAULT_PORT)
+        + "/"
+        + Iterables.getFirst(arguments.getDatabases(), "")
+        + makeIamProperties(arguments);
+  }
+
+  private static String makeScheme(String urlType) {
+    return String.format("jdbc:%s://", urlType);
+  }
+
+  private static String makeIamProperties(ConnectorArguments arguments)
+      throws UnsupportedEncodingException {
+    String profile = arguments.getIAMProfile();
+    if (profile != null) {
+      return new JdbcPropBuilder("?=&").prop("Profile", arguments.getIAMProfile()).toJdbcPart();
+    }
+    String keyId = arguments.getIAMAccessKeyID();
+    String secretKey = arguments.getIAMSecretAccessKey();
+    if (keyId != null && secretKey != null) {
+      return new JdbcPropBuilder("?=&")
+          .prop("AccessKeyID", keyId)
+          .prop("SecretAccessKey", secretKey)
+          .propOrError("DbUser", arguments.getUser(), "--user must be specified")
+          .toJdbcPart();
+    } else {
+      // Use the default IAM from ~/.aws/credentials/
+      return "";
+    }
+  }
+}
