From 7dc740a1f13fbb9c7832502ad69f1b8541cfc4c4 Mon Sep 17 00:00:00 2001 From: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> Date: Sat, 3 Oct 2026 00:15:28 +0000 Subject: [PATCH 1/5] Add typed feature flag getters Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --- .../DatabricksDriverFeatureFlagsContext.java | 74 ++++++++++++++++++- ...tabricksDriverFeatureFlagsContextTest.java | 30 +++++++- 2 files changed, 102 insertions(+), 2 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java b/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java index fcb4458414..06f0f117ce 100644 --- a/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java +++ b/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java @@ -10,11 +10,18 @@ import com.databricks.jdbc.exception.DatabricksHttpException; import com.databricks.jdbc.log.JdbcLogger; import com.databricks.jdbc.log.JdbcLoggerFactory; +import com.fasterxml.jackson.databind.JsonNode; import com.google.common.annotations.VisibleForTesting; import com.google.common.cache.Cache; import com.google.common.cache.CacheBuilder; import java.io.IOException; +import java.util.ArrayList; +import java.util.List; import java.util.Map; +import java.util.Optional; +import java.util.OptionalDouble; +import java.util.OptionalInt; +import java.util.OptionalLong; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.ScheduledFuture; @@ -153,8 +160,73 @@ void updateConnectionContext(IDatabricksConnectionContext newContext) { } public boolean isFeatureEnabled(String name) { + return getBoolean(name).orElse(false); + } + + public Optional getBoolean(String name) { + JsonNode value = parse(name); + return value != null && value.isBoolean() + ? Optional.of(value.booleanValue()) + : Optional.empty(); + } + + public OptionalInt getInt32(String name) { + JsonNode value = parse(name); + return value != null && value.isIntegralNumber() && value.canConvertToInt() + ? OptionalInt.of(value.intValue()) + : OptionalInt.empty(); + } + + public OptionalLong getInt64(String name) { + JsonNode value = parse(name); + return value != null && value.isIntegralNumber() && value.canConvertToLong() + ? OptionalLong.of(value.longValue()) + : OptionalLong.empty(); + } + + public OptionalDouble getDouble(String name) { + JsonNode value = parse(name); + if (value == null || !value.isNumber()) { + return OptionalDouble.empty(); + } + double number = value.doubleValue(); + return Double.isFinite(number) ? OptionalDouble.of(number) : OptionalDouble.empty(); + } + + public Optional getString(String name) { + JsonNode value = parse(name); + return value != null && value.isTextual() ? Optional.of(value.textValue()) : Optional.empty(); + } + + public Optional> getStringList(String name) { + JsonNode value = parse(name); + if (value == null || !value.isArray()) { + return Optional.empty(); + } + List result = new ArrayList<>(value.size()); + for (JsonNode item : value) { + if (!item.isTextual()) { + return Optional.empty(); + } + result.add(item.textValue()); + } + return Optional.of(List.copyOf(result)); + } + + private JsonNode parse(String name) { + if (name == null || name.isEmpty()) { + return null; + } String value = featureFlags.getIfPresent(name); - return Boolean.parseBoolean(value); + if (value == null) { + return null; + } + try { + return JsonUtil.getMapper().readTree(value); + } catch (IOException | RuntimeException e) { + LOGGER.debug("Feature flag {} has malformed JSON; using the consumer default", name); + return null; + } } public void shutdown() { diff --git a/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextTest.java b/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextTest.java index fa0f5b812b..eb7b956c1a 100644 --- a/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextTest.java +++ b/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextTest.java @@ -42,9 +42,12 @@ class DatabricksDriverFeatureFlagsContextTest { private DatabricksDriverFeatureFlagsContext context; @BeforeEach - void setUp() { + void setUp() throws Exception { // Mock the host for OAuth to return a test host when(connectionContextMock.getHostForOAuth()).thenReturn("test-host"); + lenient() + .when(objectMapperMock.readTree(anyString())) + .thenAnswer(invocation -> new ObjectMapper().readTree(invocation.getArgument(0))); context = new DatabricksDriverFeatureFlagsContext(connectionContextMock, new HashMap<>()); } @@ -205,6 +208,31 @@ void testIsFeatureEnabled() { assertFalse(context.isFeatureEnabled("nonexistent")); } + @Test + void testTypedGetters() { + context = + new DatabricksDriverFeatureFlagsContext( + connectionContextMock, + Map.of( + "boolean", "true", + "int32", Integer.toString(Integer.MIN_VALUE), + "int64", Long.toString(Long.MAX_VALUE), + "double", "3.5", + "string", "\"hello\"", + "string-list", "[\"a\",\"b\"]", + "wrong-type", "\"true\"", + "malformed", "not-json")); + + assertTrue(context.getBoolean("boolean").orElseThrow()); + assertEquals(Integer.MIN_VALUE, context.getInt32("int32").orElseThrow()); + assertEquals(Long.MAX_VALUE, context.getInt64("int64").orElseThrow()); + assertEquals(3.5, context.getDouble("double").orElseThrow()); + assertEquals("hello", context.getString("string").orElseThrow()); + assertEquals(List.of("a", "b"), context.getStringList("string-list").orElseThrow()); + assertTrue(context.getBoolean("wrong-type").isEmpty()); + assertTrue(context.getString("malformed").isEmpty()); + } + // ===== Additional Integration Tests ===== @Test From 6fb4af16c9ada01afebf157e13e2205ea535dcdb Mon Sep 17 00:00:00 2001 From: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> Date: Mon, 5 Oct 2026 20:17:38 +0000 Subject: [PATCH 2/5] Replace feature flag boolean wrapper with typed getter Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --- .../api/impl/DatabricksConnectionContext.java | 6 ++- .../DatabricksDriverFeatureFlagsContext.java | 4 -- .../jdbc/telemetry/TelemetryHelper.java | 3 +- ...sDriverFeatureFlagsContextFactoryTest.java | 18 +++---- ...tabricksDriverFeatureFlagsContextTest.java | 53 ++++++++++--------- 5 files changed, 43 insertions(+), 41 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java index de58419856..2ab66c01b8 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java @@ -569,7 +569,8 @@ public DatabricksClientType getClientTypeFromContext() { } // Check feature flag to determine if SEA client should be enabled if (DatabricksDriverFeatureFlagsContextFactory.getInstance(this) - .isFeatureEnabled(SQL_EXEC_FLAG_NAME)) { + .getBoolean(SQL_EXEC_FLAG_NAME) + .orElse(false)) { return DatabricksClientType.SEA; } // Default to THRIFT if feature flag is not enabled or cannot be determined @@ -1394,7 +1395,8 @@ private boolean resolveFeatureFlag(DatabricksJdbcUrlParams clientParam, String s try { serverEnabled = DatabricksDriverFeatureFlagsContextFactory.getInstance(this) - .isFeatureEnabled(serverFlagName); + .getBoolean(serverFlagName) + .orElse(false); } catch (Exception e) { LOGGER.debug("Failed to check server-side flag {}: {}", serverFlagName, e.getMessage()); } diff --git a/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java b/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java index 06f0f117ce..d2cb0c2b04 100644 --- a/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java +++ b/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java @@ -159,10 +159,6 @@ void updateConnectionContext(IDatabricksConnectionContext newContext) { this.connectionContext = newContext; } - public boolean isFeatureEnabled(String name) { - return getBoolean(name).orElse(false); - } - public Optional getBoolean(String name) { JsonNode value = parse(name); return value != null && value.isBoolean() diff --git a/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java b/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java index 4e45cfe131..e4453206de 100644 --- a/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java +++ b/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java @@ -74,7 +74,8 @@ public static boolean isTelemetryAllowedForConnection(IDatabricksConnectionConte } return context.isTelemetryEnabled() && DatabricksDriverFeatureFlagsContextFactory.getInstance(context) - .isFeatureEnabled(TELEMETRY_FEATURE_FLAG_NAME); + .getBoolean(TELEMETRY_FEATURE_FLAG_NAME) + .orElse(false); } public static void exportTelemetryLog( diff --git a/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextFactoryTest.java b/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextFactoryTest.java index 13e198722e..d449079284 100644 --- a/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextFactoryTest.java +++ b/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextFactoryTest.java @@ -192,14 +192,14 @@ void testContextPersistsUntilLastRemoval() { DatabricksDriverFeatureFlagsContext context = DatabricksDriverFeatureFlagsContextFactory.getInstance(connX); - assertTrue(context.isFeatureEnabled("test.flag")); + assertTrue(context.getBoolean("test.flag").orElse(false)); // Remove connX — context should still exist because connY is still open DatabricksDriverFeatureFlagsContextFactory.removeInstance(connX); DatabricksDriverFeatureFlagsContext contextAfterX = DatabricksDriverFeatureFlagsContextFactory.getInstance(connY); - assertTrue(contextAfterX.isFeatureEnabled("test.flag")); + assertTrue(contextAfterX.getBoolean("test.flag").orElse(false)); // Clean up DatabricksDriverFeatureFlagsContextFactory.removeInstance(connY); @@ -217,8 +217,8 @@ void testSetFeatureFlagsContextWorks() { DatabricksDriverFeatureFlagsContext context = DatabricksDriverFeatureFlagsContextFactory.getInstance(connectionContext1); - assertTrue(context.isFeatureEnabled("feature1")); - assertFalse(context.isFeatureEnabled("feature2")); + assertTrue(context.getBoolean("feature1").orElse(false)); + assertFalse(context.getBoolean("feature2").orElse(false)); } @Test @@ -251,7 +251,7 @@ void testMultipleConnectionsToSameWorkspaceShareFlags() { DatabricksDriverFeatureFlagsContextFactory.getInstance(conn2); assertSame(context1, context2); - assertTrue(context2.isFeatureEnabled("shared.flag")); + assertTrue(context2.getBoolean("shared.flag").orElse(false)); // Clean up DatabricksDriverFeatureFlagsContextFactory.removeInstance(conn1); @@ -276,10 +276,10 @@ void testDifferentWorkspacesHaveIsolatedFlags() { DatabricksDriverFeatureFlagsContextFactory.getInstance(connectionContext2); // Verify flags are isolated - assertTrue(context1.isFeatureEnabled("workspace1.flag")); - assertFalse(context1.isFeatureEnabled("workspace2.flag")); + assertTrue(context1.getBoolean("workspace1.flag").orElse(false)); + assertFalse(context1.getBoolean("workspace2.flag").orElse(false)); - assertFalse(context2.isFeatureEnabled("workspace1.flag")); - assertTrue(context2.isFeatureEnabled("workspace2.flag")); + assertFalse(context2.getBoolean("workspace1.flag").orElse(false)); + assertTrue(context2.getBoolean("workspace2.flag").orElse(false)); } } diff --git a/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextTest.java b/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextTest.java index eb7b956c1a..1e3f856691 100644 --- a/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextTest.java +++ b/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextTest.java @@ -107,7 +107,7 @@ void testFetchAndSetFlagsFromServer_Success() throws Exception { .thenReturn(response); HttpGet request = new HttpGet(FEATURE_FLAGS_ENDPOINT); context.fetchAndSetFlagsFromServer(httpClientMock, request); - assertTrue(context.isFeatureEnabled(FEATURE_FLAG_NAME)); + assertTrue(context.getBoolean(FEATURE_FLAG_NAME).orElse(false)); verify(httpClientMock).execute(request); } } @@ -129,7 +129,7 @@ void testFetchAndSetFlagsFromServer_WithCustomTTL() throws Exception { .thenReturn(response); HttpGet request = new HttpGet(FEATURE_FLAGS_ENDPOINT); context.fetchAndSetFlagsFromServer(httpClientMock, request); - assertTrue(context.isFeatureEnabled(FEATURE_FLAG_NAME)); + assertTrue(context.getBoolean(FEATURE_FLAG_NAME).orElse(false)); verify(httpClientMock).execute(request); } } @@ -141,7 +141,7 @@ void testFetchAndSetFlagsFromServer_HttpError() throws IOException, DatabricksHt when(httpClientMock.execute(any(HttpGet.class))).thenReturn(httpResponseMock); HttpGet request = new HttpGet(FEATURE_FLAGS_ENDPOINT); context.fetchAndSetFlagsFromServer(httpClientMock, request); - assertFalse(context.isFeatureEnabled(FEATURE_FLAG_NAME)); + assertFalse(context.getBoolean(FEATURE_FLAG_NAME).orElse(false)); verify(httpClientMock).execute(request); } @@ -161,7 +161,7 @@ void testFetchAndSetFlagsFromServer_EmptyFlags() throws Exception { .thenReturn(response); HttpGet request = new HttpGet(FEATURE_FLAGS_ENDPOINT); context.fetchAndSetFlagsFromServer(httpClientMock, request); - assertFalse(context.isFeatureEnabled(FEATURE_FLAG_NAME)); + assertFalse(context.getBoolean(FEATURE_FLAG_NAME).orElse(false)); verify(httpClientMock).execute(request); } } @@ -182,30 +182,30 @@ void testFetchAndSetFlagsFromServer_NullFlags() throws Exception { .thenReturn(response); HttpGet request = new HttpGet(FEATURE_FLAGS_ENDPOINT); context.fetchAndSetFlagsFromServer(httpClientMock, request); - assertFalse(context.isFeatureEnabled(FEATURE_FLAG_NAME)); + assertFalse(context.getBoolean(FEATURE_FLAG_NAME).orElse(false)); verify(httpClientMock).execute(request); } } @Test - void testIsFeatureEnabled() { + void testGetBoolean() { // Test with valid boolean values Map flags = new HashMap<>(); flags.put("flag1", "true"); flags.put("flag2", "false"); context = new DatabricksDriverFeatureFlagsContext(connectionContextMock, flags); - assertTrue(context.isFeatureEnabled("flag1")); - assertFalse(context.isFeatureEnabled("flag2")); + assertTrue(context.getBoolean("flag1").orElse(false)); + assertFalse(context.getBoolean("flag2").orElse(false)); // Test with invalid values flags.put("flag3", "invalid"); flags.put("flag4", "yes"); context = new DatabricksDriverFeatureFlagsContext(connectionContextMock, flags); - assertFalse(context.isFeatureEnabled("flag3")); - assertFalse(context.isFeatureEnabled("flag4")); + assertFalse(context.getBoolean("flag3").orElse(false)); + assertFalse(context.getBoolean("flag4").orElse(false)); // Test with non-existent flag - assertFalse(context.isFeatureEnabled("nonexistent")); + assertFalse(context.getBoolean("nonexistent").orElse(false)); } @Test @@ -236,14 +236,15 @@ void testTypedGetters() { // ===== Additional Integration Tests ===== @Test - void testIsFeatureEnabledForSqlExecFlag() { + void testGetBooleanForSqlExecFlag() { Map flags = new HashMap<>(); flags.put("databricks.partnerplatform.clientConfigsFeatureFlags.enableSqlExecForJdbc", "true"); context = new DatabricksDriverFeatureFlagsContext(connectionContextMock, flags); assertTrue( - context.isFeatureEnabled( - "databricks.partnerplatform.clientConfigsFeatureFlags.enableSqlExecForJdbc")); + context + .getBoolean("databricks.partnerplatform.clientConfigsFeatureFlags.enableSqlExecForJdbc") + .orElse(false)); } @Test @@ -272,11 +273,13 @@ void testMultipleFeatureFlagsInResponse() throws Exception { HttpGet request = new HttpGet(FEATURE_FLAGS_ENDPOINT); context.fetchAndSetFlagsFromServer(httpClientMock, request); - assertTrue(context.isFeatureEnabled("flag1")); + assertTrue(context.getBoolean("flag1").orElse(false)); assertTrue( - context.isFeatureEnabled( - "databricks.partnerplatform.clientConfigsFeatureFlags.enableSqlExecForJdbc")); - assertFalse(context.isFeatureEnabled("flag3")); + context + .getBoolean( + "databricks.partnerplatform.clientConfigsFeatureFlags.enableSqlExecForJdbc") + .orElse(false)); + assertFalse(context.getBoolean("flag3").orElse(false)); } } @@ -290,7 +293,7 @@ void testFetchAndSetFlagsFromServer_404Error() throws IOException, DatabricksHtt context.fetchAndSetFlagsFromServer(httpClientMock, request); // Should not throw, and feature should be disabled by default - assertFalse(context.isFeatureEnabled(FEATURE_FLAG_NAME)); + assertFalse(context.getBoolean(FEATURE_FLAG_NAME).orElse(false)); } @Test @@ -302,7 +305,7 @@ void testFetchAndSetFlagsFromServer_403Error() throws IOException, DatabricksHtt HttpGet request = new HttpGet(FEATURE_FLAGS_ENDPOINT); context.fetchAndSetFlagsFromServer(httpClientMock, request); - assertFalse(context.isFeatureEnabled(FEATURE_FLAG_NAME)); + assertFalse(context.getBoolean(FEATURE_FLAG_NAME).orElse(false)); } @Test @@ -314,17 +317,17 @@ void testFetchAndSetFlagsFromServer_503Error() throws IOException, DatabricksHtt HttpGet request = new HttpGet(FEATURE_FLAGS_ENDPOINT); context.fetchAndSetFlagsFromServer(httpClientMock, request); - assertFalse(context.isFeatureEnabled(FEATURE_FLAG_NAME)); + assertFalse(context.getBoolean(FEATURE_FLAG_NAME).orElse(false)); } @Test - void testIsFeatureEnabledCaseSensitive() { + void testGetBooleanCaseSensitive() { Map flags = new HashMap<>(); flags.put("TestFlag", "true"); context = new DatabricksDriverFeatureFlagsContext(connectionContextMock, flags); - assertTrue(context.isFeatureEnabled("TestFlag")); - assertFalse(context.isFeatureEnabled("testflag")); - assertFalse(context.isFeatureEnabled("TESTFLAG")); + assertTrue(context.getBoolean("TestFlag").orElse(false)); + assertFalse(context.getBoolean("testflag").orElse(false)); + assertFalse(context.getBoolean("TESTFLAG").orElse(false)); } } From 7cd97a15d04bf06fd5b68882f97b8bea617f936a Mon Sep 17 00:00:00 2001 From: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> Date: Mon, 5 Oct 2026 20:19:22 +0000 Subject: [PATCH 3/5] Document typed SAFE feature flag getters Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --- NEXT_CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index ac3a5cdda0..4252c36ea6 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -4,6 +4,8 @@ ### Added +- Added typed getters for all six SAFE feature flag types using the existing cache. + ### Updated - Bumped Apache HttpClient 5 (`httpclient5`) from 5.6.3 to 5.6.4. From 7f4a6e5d4f261012ee58e3244a6df9f012aa5212 Mon Sep 17 00:00:00 2001 From: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> Date: Mon, 5 Oct 2026 20:39:18 +0000 Subject: [PATCH 4/5] Preserve boolean feature flag defaults in getBoolean Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --- .../api/impl/DatabricksConnectionContext.java | 7 +- .../DatabricksDriverFeatureFlagsContext.java | 9 ++- .../jdbc/telemetry/TelemetryHelper.java | 3 +- ...sDriverFeatureFlagsContextFactoryTest.java | 18 ++--- ...tabricksDriverFeatureFlagsContextTest.java | 66 +++++++++---------- 5 files changed, 49 insertions(+), 54 deletions(-) diff --git a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java index 2ab66c01b8..639f4d23b9 100644 --- a/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java +++ b/src/main/java/com/databricks/jdbc/api/impl/DatabricksConnectionContext.java @@ -569,8 +569,7 @@ public DatabricksClientType getClientTypeFromContext() { } // Check feature flag to determine if SEA client should be enabled if (DatabricksDriverFeatureFlagsContextFactory.getInstance(this) - .getBoolean(SQL_EXEC_FLAG_NAME) - .orElse(false)) { + .getBoolean(SQL_EXEC_FLAG_NAME)) { return DatabricksClientType.SEA; } // Default to THRIFT if feature flag is not enabled or cannot be determined @@ -1394,9 +1393,7 @@ private boolean resolveFeatureFlag(DatabricksJdbcUrlParams clientParam, String s boolean serverEnabled = false; try { serverEnabled = - DatabricksDriverFeatureFlagsContextFactory.getInstance(this) - .getBoolean(serverFlagName) - .orElse(false); + DatabricksDriverFeatureFlagsContextFactory.getInstance(this).getBoolean(serverFlagName); } catch (Exception e) { LOGGER.debug("Failed to check server-side flag {}: {}", serverFlagName, e.getMessage()); } diff --git a/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java b/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java index d2cb0c2b04..4364ff9a6e 100644 --- a/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java +++ b/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java @@ -159,11 +159,10 @@ void updateConnectionContext(IDatabricksConnectionContext newContext) { this.connectionContext = newContext; } - public Optional getBoolean(String name) { - JsonNode value = parse(name); - return value != null && value.isBoolean() - ? Optional.of(value.booleanValue()) - : Optional.empty(); + /** Returns true only for a case-insensitive "true" value; missing values default to false. */ + public boolean getBoolean(String name) { + String value = featureFlags.getIfPresent(name); + return Boolean.parseBoolean(value); } public OptionalInt getInt32(String name) { diff --git a/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java b/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java index e4453206de..45ba228cb8 100644 --- a/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java +++ b/src/main/java/com/databricks/jdbc/telemetry/TelemetryHelper.java @@ -74,8 +74,7 @@ public static boolean isTelemetryAllowedForConnection(IDatabricksConnectionConte } return context.isTelemetryEnabled() && DatabricksDriverFeatureFlagsContextFactory.getInstance(context) - .getBoolean(TELEMETRY_FEATURE_FLAG_NAME) - .orElse(false); + .getBoolean(TELEMETRY_FEATURE_FLAG_NAME); } public static void exportTelemetryLog( diff --git a/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextFactoryTest.java b/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextFactoryTest.java index d449079284..117b5ff03e 100644 --- a/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextFactoryTest.java +++ b/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextFactoryTest.java @@ -192,14 +192,14 @@ void testContextPersistsUntilLastRemoval() { DatabricksDriverFeatureFlagsContext context = DatabricksDriverFeatureFlagsContextFactory.getInstance(connX); - assertTrue(context.getBoolean("test.flag").orElse(false)); + assertTrue(context.getBoolean("test.flag")); // Remove connX — context should still exist because connY is still open DatabricksDriverFeatureFlagsContextFactory.removeInstance(connX); DatabricksDriverFeatureFlagsContext contextAfterX = DatabricksDriverFeatureFlagsContextFactory.getInstance(connY); - assertTrue(contextAfterX.getBoolean("test.flag").orElse(false)); + assertTrue(contextAfterX.getBoolean("test.flag")); // Clean up DatabricksDriverFeatureFlagsContextFactory.removeInstance(connY); @@ -217,8 +217,8 @@ void testSetFeatureFlagsContextWorks() { DatabricksDriverFeatureFlagsContext context = DatabricksDriverFeatureFlagsContextFactory.getInstance(connectionContext1); - assertTrue(context.getBoolean("feature1").orElse(false)); - assertFalse(context.getBoolean("feature2").orElse(false)); + assertTrue(context.getBoolean("feature1")); + assertFalse(context.getBoolean("feature2")); } @Test @@ -251,7 +251,7 @@ void testMultipleConnectionsToSameWorkspaceShareFlags() { DatabricksDriverFeatureFlagsContextFactory.getInstance(conn2); assertSame(context1, context2); - assertTrue(context2.getBoolean("shared.flag").orElse(false)); + assertTrue(context2.getBoolean("shared.flag")); // Clean up DatabricksDriverFeatureFlagsContextFactory.removeInstance(conn1); @@ -276,10 +276,10 @@ void testDifferentWorkspacesHaveIsolatedFlags() { DatabricksDriverFeatureFlagsContextFactory.getInstance(connectionContext2); // Verify flags are isolated - assertTrue(context1.getBoolean("workspace1.flag").orElse(false)); - assertFalse(context1.getBoolean("workspace2.flag").orElse(false)); + assertTrue(context1.getBoolean("workspace1.flag")); + assertFalse(context1.getBoolean("workspace2.flag")); - assertFalse(context2.getBoolean("workspace1.flag").orElse(false)); - assertTrue(context2.getBoolean("workspace2.flag").orElse(false)); + assertFalse(context2.getBoolean("workspace1.flag")); + assertTrue(context2.getBoolean("workspace2.flag")); } } diff --git a/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextTest.java b/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextTest.java index 1e3f856691..6cbe12154e 100644 --- a/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextTest.java +++ b/src/test/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContextTest.java @@ -42,12 +42,9 @@ class DatabricksDriverFeatureFlagsContextTest { private DatabricksDriverFeatureFlagsContext context; @BeforeEach - void setUp() throws Exception { + void setUp() { // Mock the host for OAuth to return a test host when(connectionContextMock.getHostForOAuth()).thenReturn("test-host"); - lenient() - .when(objectMapperMock.readTree(anyString())) - .thenAnswer(invocation -> new ObjectMapper().readTree(invocation.getArgument(0))); context = new DatabricksDriverFeatureFlagsContext(connectionContextMock, new HashMap<>()); } @@ -107,7 +104,7 @@ void testFetchAndSetFlagsFromServer_Success() throws Exception { .thenReturn(response); HttpGet request = new HttpGet(FEATURE_FLAGS_ENDPOINT); context.fetchAndSetFlagsFromServer(httpClientMock, request); - assertTrue(context.getBoolean(FEATURE_FLAG_NAME).orElse(false)); + assertTrue(context.getBoolean(FEATURE_FLAG_NAME)); verify(httpClientMock).execute(request); } } @@ -129,7 +126,7 @@ void testFetchAndSetFlagsFromServer_WithCustomTTL() throws Exception { .thenReturn(response); HttpGet request = new HttpGet(FEATURE_FLAGS_ENDPOINT); context.fetchAndSetFlagsFromServer(httpClientMock, request); - assertTrue(context.getBoolean(FEATURE_FLAG_NAME).orElse(false)); + assertTrue(context.getBoolean(FEATURE_FLAG_NAME)); verify(httpClientMock).execute(request); } } @@ -141,7 +138,7 @@ void testFetchAndSetFlagsFromServer_HttpError() throws IOException, DatabricksHt when(httpClientMock.execute(any(HttpGet.class))).thenReturn(httpResponseMock); HttpGet request = new HttpGet(FEATURE_FLAGS_ENDPOINT); context.fetchAndSetFlagsFromServer(httpClientMock, request); - assertFalse(context.getBoolean(FEATURE_FLAG_NAME).orElse(false)); + assertFalse(context.getBoolean(FEATURE_FLAG_NAME)); verify(httpClientMock).execute(request); } @@ -161,7 +158,7 @@ void testFetchAndSetFlagsFromServer_EmptyFlags() throws Exception { .thenReturn(response); HttpGet request = new HttpGet(FEATURE_FLAGS_ENDPOINT); context.fetchAndSetFlagsFromServer(httpClientMock, request); - assertFalse(context.getBoolean(FEATURE_FLAG_NAME).orElse(false)); + assertFalse(context.getBoolean(FEATURE_FLAG_NAME)); verify(httpClientMock).execute(request); } } @@ -182,7 +179,7 @@ void testFetchAndSetFlagsFromServer_NullFlags() throws Exception { .thenReturn(response); HttpGet request = new HttpGet(FEATURE_FLAGS_ENDPOINT); context.fetchAndSetFlagsFromServer(httpClientMock, request); - assertFalse(context.getBoolean(FEATURE_FLAG_NAME).orElse(false)); + assertFalse(context.getBoolean(FEATURE_FLAG_NAME)); verify(httpClientMock).execute(request); } } @@ -193,19 +190,25 @@ void testGetBoolean() { Map flags = new HashMap<>(); flags.put("flag1", "true"); flags.put("flag2", "false"); + flags.put("mixedCase", "TrUe"); context = new DatabricksDriverFeatureFlagsContext(connectionContextMock, flags); - assertTrue(context.getBoolean("flag1").orElse(false)); - assertFalse(context.getBoolean("flag2").orElse(false)); + assertTrue(context.getBoolean("flag1")); + assertFalse(context.getBoolean("flag2")); + assertTrue(context.getBoolean("mixedCase")); // Test with invalid values flags.put("flag3", "invalid"); flags.put("flag4", "yes"); + flags.put("flag5", "null"); + flags.put("flag6", "\"true\""); + flags.put("flag7", ""); + flags.put("flag8", " true "); context = new DatabricksDriverFeatureFlagsContext(connectionContextMock, flags); - assertFalse(context.getBoolean("flag3").orElse(false)); - assertFalse(context.getBoolean("flag4").orElse(false)); - - // Test with non-existent flag - assertFalse(context.getBoolean("nonexistent").orElse(false)); + for (String name : + new String[] {"flag3", "flag4", "flag5", "flag6", "flag7", "flag8", "nonexistent", ""}) { + assertFalse(context.getBoolean(name)); + } + assertThrows(NullPointerException.class, () -> context.getBoolean(null)); } @Test @@ -223,13 +226,13 @@ void testTypedGetters() { "wrong-type", "\"true\"", "malformed", "not-json")); - assertTrue(context.getBoolean("boolean").orElseThrow()); + assertTrue(context.getBoolean("boolean")); assertEquals(Integer.MIN_VALUE, context.getInt32("int32").orElseThrow()); assertEquals(Long.MAX_VALUE, context.getInt64("int64").orElseThrow()); assertEquals(3.5, context.getDouble("double").orElseThrow()); assertEquals("hello", context.getString("string").orElseThrow()); assertEquals(List.of("a", "b"), context.getStringList("string-list").orElseThrow()); - assertTrue(context.getBoolean("wrong-type").isEmpty()); + assertFalse(context.getBoolean("wrong-type")); assertTrue(context.getString("malformed").isEmpty()); } @@ -242,9 +245,8 @@ void testGetBooleanForSqlExecFlag() { context = new DatabricksDriverFeatureFlagsContext(connectionContextMock, flags); assertTrue( - context - .getBoolean("databricks.partnerplatform.clientConfigsFeatureFlags.enableSqlExecForJdbc") - .orElse(false)); + context.getBoolean( + "databricks.partnerplatform.clientConfigsFeatureFlags.enableSqlExecForJdbc")); } @Test @@ -273,13 +275,11 @@ void testMultipleFeatureFlagsInResponse() throws Exception { HttpGet request = new HttpGet(FEATURE_FLAGS_ENDPOINT); context.fetchAndSetFlagsFromServer(httpClientMock, request); - assertTrue(context.getBoolean("flag1").orElse(false)); + assertTrue(context.getBoolean("flag1")); assertTrue( - context - .getBoolean( - "databricks.partnerplatform.clientConfigsFeatureFlags.enableSqlExecForJdbc") - .orElse(false)); - assertFalse(context.getBoolean("flag3").orElse(false)); + context.getBoolean( + "databricks.partnerplatform.clientConfigsFeatureFlags.enableSqlExecForJdbc")); + assertFalse(context.getBoolean("flag3")); } } @@ -293,7 +293,7 @@ void testFetchAndSetFlagsFromServer_404Error() throws IOException, DatabricksHtt context.fetchAndSetFlagsFromServer(httpClientMock, request); // Should not throw, and feature should be disabled by default - assertFalse(context.getBoolean(FEATURE_FLAG_NAME).orElse(false)); + assertFalse(context.getBoolean(FEATURE_FLAG_NAME)); } @Test @@ -305,7 +305,7 @@ void testFetchAndSetFlagsFromServer_403Error() throws IOException, DatabricksHtt HttpGet request = new HttpGet(FEATURE_FLAGS_ENDPOINT); context.fetchAndSetFlagsFromServer(httpClientMock, request); - assertFalse(context.getBoolean(FEATURE_FLAG_NAME).orElse(false)); + assertFalse(context.getBoolean(FEATURE_FLAG_NAME)); } @Test @@ -317,7 +317,7 @@ void testFetchAndSetFlagsFromServer_503Error() throws IOException, DatabricksHtt HttpGet request = new HttpGet(FEATURE_FLAGS_ENDPOINT); context.fetchAndSetFlagsFromServer(httpClientMock, request); - assertFalse(context.getBoolean(FEATURE_FLAG_NAME).orElse(false)); + assertFalse(context.getBoolean(FEATURE_FLAG_NAME)); } @Test @@ -326,8 +326,8 @@ void testGetBooleanCaseSensitive() { flags.put("TestFlag", "true"); context = new DatabricksDriverFeatureFlagsContext(connectionContextMock, flags); - assertTrue(context.getBoolean("TestFlag").orElse(false)); - assertFalse(context.getBoolean("testflag").orElse(false)); - assertFalse(context.getBoolean("TESTFLAG").orElse(false)); + assertTrue(context.getBoolean("TestFlag")); + assertFalse(context.getBoolean("testflag")); + assertFalse(context.getBoolean("TESTFLAG")); } } From a64229427f959c9cdafd24b396a83805d3c0bdcd Mon Sep 17 00:00:00 2001 From: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> Date: Mon, 5 Oct 2026 23:19:33 +0000 Subject: [PATCH 5/5] Avoid redundant feature flag string-list copy Signed-off-by: Cathleen Yan <58714163+cathleeny@users.noreply.github.com> --- .../jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java b/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java index 4364ff9a6e..cc1cf63c55 100644 --- a/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java +++ b/src/main/java/com/databricks/jdbc/common/safe/DatabricksDriverFeatureFlagsContext.java @@ -16,6 +16,7 @@ import com.google.common.cache.CacheBuilder; import java.io.IOException; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.Optional; @@ -205,7 +206,7 @@ public Optional> getStringList(String name) { } result.add(item.textValue()); } - return Optional.of(List.copyOf(result)); + return Optional.of(Collections.unmodifiableList(result)); } private JsonNode parse(String name) {