From 6e2de0f3707d95f5688897fea122257a19fc7e4e Mon Sep 17 00:00:00 2001
From: sfc-gh-ext-simba-lb <lbarber@simba.com>
Date: Thu, 6 Jul 2023 12:41:39 -0700
Subject: [PATCH] allow connection property to disable OOB telemetry

---
 .../net/snowflake/client/core/SFSession.java  |  9 ++++
 .../jdbc/telemetryOOB/TelemetryServiceIT.java | 54 ++++++++++++++++++-
 2 files changed, 62 insertions(+), 1 deletion(-)

diff --git a/src/main/java/net/snowflake/client/core/SFSession.java b/src/main/java/net/snowflake/client/core/SFSession.java
index a6bd5e375..0a2fee837 100644
--- a/src/main/java/net/snowflake/client/core/SFSession.java
+++ b/src/main/java/net/snowflake/client/core/SFSession.java
@@ -484,6 +484,15 @@ public synchronized void open() throws SFException, SnowflakeSQLException {
         .setOCSPMode(getOCSPMode())
         .setHttpClientSettingsKey(httpClientSettingsKey);
 
+    // Enable or disable OOB telemetry based on connection parameter. Default is disabled.
+    // The value may still change later when session parameters from the server are read.
+    if (getBooleanValue(
+        connectionPropertiesMap.get(SFSessionProperty.CLIENT_OUT_OF_BAND_TELEMETRY_ENABLED))) {
+      TelemetryService.enable();
+    } else {
+      TelemetryService.disable();
+    }
+
     // propagate OCSP mode to SFTrustManager. Note OCSP setting is global on JVM.
     HttpUtil.initHttpClient(httpClientSettingsKey, null);
     SFLoginOutput loginOutput =
diff --git a/src/test/java/net/snowflake/client/jdbc/telemetryOOB/TelemetryServiceIT.java b/src/test/java/net/snowflake/client/jdbc/telemetryOOB/TelemetryServiceIT.java
index 5d7df7a97..1850a6732 100644
--- a/src/test/java/net/snowflake/client/jdbc/telemetryOOB/TelemetryServiceIT.java
+++ b/src/test/java/net/snowflake/client/jdbc/telemetryOOB/TelemetryServiceIT.java
@@ -2,7 +2,7 @@
 
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.MatcherAssert.assertThat;
-import static org.junit.Assert.fail;
+import static org.junit.Assert.*;
 
 import java.sql.Connection;
 import java.sql.SQLException;
@@ -308,6 +308,58 @@ public void testHTAPTelemetry() throws SQLException {
     }
   }
 
+  /**
+   * Requires part 2 of SNOW-844477. Make sure CLIENT_OUT_OF_BAND_TELEMETRY_ENABLED is true at
+   * account level. Tests connection property CLIENT_OUT_OF_BAND_TELEMETRY_ENABLED=true
+   */
+  @Ignore
+  @Test
+  public void testOOBTelemetryEnabled() throws SQLException {
+    Properties properties = new Properties();
+    properties.put("CLIENT_OUT_OF_BAND_TELEMETRY_ENABLED", "true");
+    try (Connection con = getConnection(properties)) {
+      Statement statement = con.createStatement();
+      statement.execute("select 1");
+      // Make sure OOB telemetry is enabled
+      assertTrue(TelemetryService.getInstance().isEnabled());
+    }
+  }
+
+  /**
+   * Requires part 2 of SNOW-844477. Make sure CLIENT_OUT_OF_BAND_TELEMETRY_ENABLED is false at
+   * account level. Tests connection property CLIENT_OUT_OF_BAND_TELEMETRY_ENABLED=false
+   */
+  @Ignore
+  @Test
+  public void testOOBTelemetryDisabled() throws SQLException {
+    Properties properties = new Properties();
+    properties.put("CLIENT_OUT_OF_BAND_TELEMETRY_ENABLED", "false");
+    try (Connection con = getConnection(properties)) {
+      Statement statement = con.createStatement();
+      statement.execute("select 1");
+      // Make sure OOB telemetry is disabled
+      assertFalse(TelemetryService.getInstance().isEnabled());
+    }
+  }
+
+  /**
+   * Requires part 2 of SNOW-844477. Make sure CLIENT_OUT_OF_BAND_TELEMETRY_ENABLED is true at
+   * account level. Tests connection property CLIENT_OUT_OF_BAND_TELEMETRY_ENABLED=false but
+   * CLIENT_OUT_OF_BAND_TELEMETRY_ENABLED is enabled on account level
+   */
+  @Ignore
+  @Test
+  public void testOOBTelemetryEnabledOnServerDisabledOnClient() throws SQLException {
+    Properties properties = new Properties();
+    properties.put("CLIENT_OUT_OF_BAND_TELEMETRY_ENABLED", "false");
+    try (Connection con = getConnection(properties)) {
+      Statement statement = con.createStatement();
+      statement.execute("select 1");
+      // Make sure OOB telemetry is enabled
+      assertTrue(TelemetryService.getInstance().isEnabled());
+    }
+  }
+
   /**
    * Test case for checking telemetry message for SnowflakeSQLExceptions. In-band telemetry should
    * be used.