From 4a85222fb7d53e0f112ff7c885a31b6cdc3a1446 Mon Sep 17 00:00:00 2001
From: Vikrant Puppala <vikrant.puppala@databricks.com>
Date: Wed, 19 Nov 2025 20:29:58 +0530
Subject: [PATCH] Improve logging when jdbc is shaded

Signed-off-by: Vikrant Puppala <vikrant.puppala@databricks.com>
---
 .../com/databricks/jdbc/log/JulLogger.java    | 24 +++++++
 .../databricks/jdbc/log/JulLoggerTest.java    | 70 +++++++++++++++++++
 2 files changed, 94 insertions(+)

diff --git a/src/main/java/com/databricks/jdbc/log/JulLogger.java b/src/main/java/com/databricks/jdbc/log/JulLogger.java
index f984812a8..aa620d8d4 100644
--- a/src/main/java/com/databricks/jdbc/log/JulLogger.java
+++ b/src/main/java/com/databricks/jdbc/log/JulLogger.java
@@ -233,6 +233,19 @@ private static String getPackagePrefix() {
     if (prefix != null && !prefix.isEmpty()) {
       return prefix;
     }
+    // Auto-detect the actual package prefix by using the current class
+    // This handles shaded JARs where the package might be relocated
+    String actualPackageName = JulLogger.class.getPackage().getName();
+    // Extract the root prefix (e.g., "jdbc.shaded.v1.0.12.OSS.com.databricks.jdbc"
+    // or "com.databricks.jdbc")
+    if (actualPackageName.contains(DEFAULT_PACKAGE_PREFIX)) {
+      // Find the actual prefix including any shading prefix
+      int defaultPrefixIndex = actualPackageName.indexOf(DEFAULT_PACKAGE_PREFIX);
+      if (defaultPrefixIndex > 0) {
+        // Include everything up to and including the default prefix
+        return actualPackageName.substring(0, defaultPrefixIndex + DEFAULT_PACKAGE_PREFIX.length());
+      }
+    }
     return DEFAULT_PACKAGE_PREFIX;
   }
 
@@ -241,6 +254,17 @@ private static String getDriverPackagePrefix() {
     if (StringUtils.isNotEmpty(prefix)) {
       return prefix;
     }
+    // Auto-detect the actual driver package prefix
+    // Try to determine if the driver package is also shaded
+    String mainPackagePrefix = getPackagePrefix();
+    // Check if we have a shaded prefix
+    if (!mainPackagePrefix.equals(DEFAULT_PACKAGE_PREFIX)
+        && mainPackagePrefix.contains(DEFAULT_PACKAGE_PREFIX)) {
+      // Extract the shading prefix and apply it to the driver package
+      String shadingPrefix =
+          mainPackagePrefix.substring(0, mainPackagePrefix.indexOf(DEFAULT_PACKAGE_PREFIX));
+      return shadingPrefix + DEFAULT_DRIVER_PACKAGE_PREFIX;
+    }
     return DEFAULT_DRIVER_PACKAGE_PREFIX;
   }
 
diff --git a/src/test/java/com/databricks/jdbc/log/JulLoggerTest.java b/src/test/java/com/databricks/jdbc/log/JulLoggerTest.java
index 3533263f1..43b366843 100644
--- a/src/test/java/com/databricks/jdbc/log/JulLoggerTest.java
+++ b/src/test/java/com/databricks/jdbc/log/JulLoggerTest.java
@@ -183,4 +183,74 @@ private String[] methodCallingLogger() {
   private String[] info() {
     return JulLogger.getCaller();
   }
+
+  @Test
+  void testPackagePrefixDetectionForUnshadedJar() {
+    // Test that package prefix detection works for unshaded JARs
+    String packagePrefix = JulLogger.PARENT_CLASS_PREFIX;
+    // For unshaded JARs, it should be the default prefix
+    assertTrue(
+        packagePrefix.endsWith("com.databricks.jdbc"),
+        "Package prefix should end with com.databricks.jdbc, got: " + packagePrefix);
+  }
+
+  @Test
+  void testPackagePrefixDetectionForShadedJar() {
+    // Test that the package prefix includes any shading prefix
+    // The JulLogger should be able to detect its own package correctly
+    String actualPackageName = JulLogger.class.getPackage().getName();
+    String packagePrefix = JulLogger.PARENT_CLASS_PREFIX;
+
+    // Verify that the package prefix matches the actual runtime package structure
+    assertTrue(
+        actualPackageName.startsWith(packagePrefix),
+        String.format(
+            "Actual package '%s' should start with detected prefix '%s'",
+            actualPackageName, packagePrefix));
+
+    // If the jar is shaded, the prefix should include the shading prefix
+    if (actualPackageName.contains(".") && !actualPackageName.startsWith("com.databricks.jdbc")) {
+      // This means we have a shaded jar
+      assertTrue(
+          packagePrefix.contains("com.databricks.jdbc"),
+          "Shaded package prefix should still contain com.databricks.jdbc");
+    }
+  }
+
+  @Test
+  void testInitLoggerConfiguresCorrectPackageForShadedJar() throws IOException {
+    // This test verifies that when initLogger is called, it configures loggers
+    // for the correct package prefix (whether shaded or unshaded)
+    JulLogger.initLogger(Level.OFF, JulLogger.STDOUT, 1024, 1);
+
+    // Get the configured logger
+    Logger jdbcLogger = Logger.getLogger(JulLogger.PARENT_CLASS_PREFIX);
+
+    // Verify it was configured with the correct level
+    assertEquals(
+        Level.OFF,
+        jdbcLogger.getLevel(),
+        "Logger should be configured with Level.OFF to suppress all logs");
+
+    // Verify that a logger created with the actual package name inherits the settings
+    String testLoggerName =
+        JulLogger.class.getPackage().getName() + ".api.impl.ExecutionResultFactory";
+    Logger testLogger = Logger.getLogger(testLoggerName);
+
+    // The effective level should be OFF (inherited from parent)
+    Level effectiveLevel = testLogger.getLevel();
+    // If the local level is null, it inherits from parent
+    if (effectiveLevel == null) {
+      Logger parent = testLogger.getParent();
+      while (parent != null && effectiveLevel == null) {
+        effectiveLevel = parent.getLevel();
+        parent = parent.getParent();
+      }
+    }
+
+    assertEquals(
+        Level.OFF,
+        effectiveLevel,
+        "Child loggers should inherit Level.OFF from properly configured parent logger");
+  }
 }