SNOW-2117067 Fix prober Dockerfile (#2223)

Author: sfc-gh-rkowalski

prober/Dockerfile

@@ -1,4 +1,4 @@
-FROM ubuntu:24.04
+FROM ubuntu:25.10
 
 # boilerplate labels required by validation when pushing to ACR, ECR & GCR
 LABEL org.opencontainers.image.source="https://github.com/snowflakedb/snowflake-jdbc"
@@ -21,7 +21,7 @@ RUN apt-get update && \
     ca-certificates \
     && rm -rf /var/lib/apt/lists/*
 
-ENV SDKMAN_DIR="/root/.sdkman"
+ENV SDKMAN_DIR="/app/.sdkman"
 
 RUN curl -s "https://get.sdkman.io?rcupdate=false&ci=false" | bash
 
@@ -44,9 +44,6 @@ RUN \
 
 WORKDIR /app
 
-# Create the package directory structure for the Prober
-RUN mkdir -p com/snowflake/client/jdbc/prober
-
 # Copy the Prober.java source code into the correct package structure
 COPY /src/main/java/com/snowflake/client/jdbc/prober/Prober.java com/snowflake/client/jdbc/prober/Prober.java
 
@@ -67,8 +64,6 @@ RUN \
     bash -c "${COMPILE_COMMAND}"
 
 # Copy the entrypoint script into the image and make it executable
-COPY entrypoint.sh /entrypoint.sh
-RUN chmod +x /entrypoint.sh
-
-# Set the entrypoint script to be executed when the container starts
-ENTRYPOINT ["/entrypoint.sh"]
+COPY entrypoint.sh entrypoint.sh
+RUN chmod +x entrypoint.sh
+RUN chmod +x /app/.sdkman/bin/sdkman-init.sh

prober/entrypoint.sh

@@ -3,7 +3,6 @@
 # Exit immediately if a command exits with a non-zero status
 set -e
 
-# Assign arguments to variables for clarity
 JAVA_VERSION_ID=""
 JDBC_VERSION=""
 
@@ -17,23 +16,25 @@ else
   exit 1
 fi
 
+PROBER_ARGS=("${@}")
+
 # Define key directories
-SDKMAN_DIR="/root/.sdkman"
+SDKMAN_DIR="/app/.sdkman"
 DRIVERS_DIR="/opt/jdbc_drivers"
 PROBER_APP_DIR="/app"
 
-# Get all arguments starting from the third one to pass to the Java application
-PROBER_ARGS=("${@}")
+source "${SDKMAN_DIR}/bin/sdkman-init.sh" >/dev/null 2>&1
+sdk use java "${JAVA_VERSION_ID}" >/dev/null 2>&1
 
 JDBC_DRIVER_JAR_PATH="${DRIVERS_DIR}/snowflake-jdbc-${JDBC_VERSION}.jar"
 CURRENT_CLASSPATH="${PROBER_APP_DIR}:${JDBC_DRIVER_JAR_PATH}"
 
-# Construct the command to run the pre-compiled Prober application
-RUN_PROBER_COMMAND="source \"${SDKMAN_DIR}/bin/sdkman-init.sh\" >/dev/null 2>&1 && \
-                    sdk use java \"${JAVA_VERSION_ID}\" >/dev/null 2>&1 && \
-                    java -cp \"${CURRENT_CLASSPATH}\" --add-opens=java.base/java.nio=ALL-UNNAMED com.snowflake.client.jdbc.prober.Prober \$@"
-
-# Execute the Prober
-bash -c "${RUN_PROBER_COMMAND}" _ "${PROBER_ARGS[@]}"
+exec java -cp "${CURRENT_CLASSPATH}" \
+          --add-opens=java.base/java.nio=ALL-UNNAMED \
+          com.snowflake.client.jdbc.prober.Prober \
+          "${PROBER_ARGS[@]}"
 
-exit 0
+# This line will not be reached because 'exec' replaces the process.
+# It would only be reached if 'exec java' itself failed to start the Java process.
+echo "Error: Java application failed to start."
+exit 1 # Exit with an error if Java failed to start

prober/src/main/java/com/snowflake/client/jdbc/prober/Prober.java

@@ -8,12 +8,20 @@
 import java.io.InputStream;
 import java.io.InputStreamReader;
 import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Paths;
+import java.security.KeyFactory;
+import java.security.NoSuchAlgorithmException;
+import java.security.interfaces.RSAPrivateCrtKey;
+import java.security.spec.InvalidKeySpecException;
+import java.security.spec.PKCS8EncodedKeySpec;
 import java.sql.Connection;
 import java.sql.DriverManager;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.ArrayList;
+import java.util.Base64;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -52,20 +60,21 @@ enum Scope {
     PUT_FETCH_GET
   }
 
-  public static void main(String[] args) throws IOException, SQLException {
-    disableLogging();
+  public static void main(String[] args) throws Exception {
     Map<String, String> arguments = parseArguments(args);
 
-    String url = "jdbc:snowflake://" + arguments.get("account") + "." + arguments.get("host");
+    String url = "jdbc:snowflake://" + arguments.get("host");
     Properties props = new Properties();
     for (Map.Entry<String, String> entry : arguments.entrySet()) {
       props.setProperty(entry.getKey(), entry.getValue());
     }
+    setPrivateKey(props);
+    setupLogging(props);
 
     javaVersion = props.getProperty("java_version");
     driverVersion = props.getProperty("driver_version");
 
-    if (Scope.LOGIN.name().toLowerCase().equals(props.getProperty("scope"))){
+    if (Scope.LOGIN.name().toLowerCase().equals(props.getProperty("scope"))) {
       testLogin(url, props);
     }
     if (Scope.PUT_FETCH_GET.name().toLowerCase().equals(props.getProperty("scope"))) {
@@ -84,10 +93,10 @@ private static void testLogin(String url, Properties properties) {
     } catch (SQLException e) {
       success = false;
       System.err.println(e.getMessage());
-      logMetric("cloudprober_driver_jdbc_perform_login", Status.FAILURE);
+      logMetric("cloudprober_driver_java_perform_login", Status.FAILURE);
       System.exit(1);
     }
-    logMetric("cloudprober_driver_jdbc_perform_login", success ? Status.SUCCESS : Status.FAILURE);
+    logMetric("cloudprober_driver_java_perform_login", success ? Status.SUCCESS : Status.FAILURE);
   }
 
   private static void testPutFetchGet(String url, Properties properties) {
@@ -271,8 +280,9 @@ private static void createDataStage(Statement statement) throws SQLException {
     }
   }
 
-  private static void disableLogging() throws IOException {
-    String loggingPropertiesString = ".level=OFF";
+  private static void setupLogging(Properties properties) throws IOException {
+    String loggingPropertiesString = "handlers=java.util.logging.ConsoleHandler\n.level=" + properties.getProperty("log_level");
+    properties.put("JAVA_LOGGING_CONSOLE_STD_OUT", "false");
     try (InputStream propertiesStream = new ByteArrayInputStream(
         loggingPropertiesString.getBytes(StandardCharsets.UTF_8)
     )) {
@@ -281,7 +291,7 @@ private static void disableLogging() throws IOException {
   }
 
   private static void logMetric(String metricName, Status status) {
-    System.out.println(metricName + "{{java_version=" + javaVersion + ", driver_version=" + driverVersion + "} " + status.getCode() + "}");
+    System.out.println(metricName + "{java_version=\"" + javaVersion + "\", driver_version=\"" + driverVersion + "\"} " + status.getCode());
   }
 
   private static Map<String, String> parseArguments(String[] args) {
@@ -341,4 +351,17 @@ private static String generateRandomString(int length) {
     builder.setCharAt(0, Character.toUpperCase(builder.charAt(0)));
     return builder.toString();
   }
+
+  private static void setPrivateKey(Properties props) throws IOException, NoSuchAlgorithmException, InvalidKeySpecException {
+    String keyStr = new String(Files.readAllBytes(Paths.get(props.getProperty("private_key_file"))), StandardCharsets.UTF_8).trim();
+    byte[] keyBytes = Base64.getUrlDecoder().decode(keyStr);
+
+    // Convert the DER bytes to a private key object
+    PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(keyBytes);
+    KeyFactory keyFactory = KeyFactory.getInstance("RSA");
+    RSAPrivateCrtKey privateKey = (RSAPrivateCrtKey) keyFactory.generatePrivate(keySpec);
+    props.put("privateKey", privateKey);
+    // Remove the path from properties so the driver does not try to read it
+    props.remove("private_key_file");
+  }
 }