diff --git a/src/main/java/net/snowflake/client/core/SFBaseSession.java b/src/main/java/net/snowflake/client/core/SFBaseSession.java index db5cd445e..b440f5fcc 100644 --- a/src/main/java/net/snowflake/client/core/SFBaseSession.java +++ b/src/main/java/net/snowflake/client/core/SFBaseSession.java @@ -74,6 +74,7 @@ public abstract class SFBaseSession { private boolean formatDateWithTimezone; private boolean enableCombineDescribe; private boolean clientTelemetryEnabled = false; + private boolean isTelemetryServiceAvailable = true; private boolean useSessionTimezone; private boolean defaultFormatDateWithTimezone = true; private boolean getDateUseNullTimezone = true; @@ -1012,6 +1013,14 @@ public void setClientTelemetryEnabled(boolean clientTelemetryEnabled) { this.clientTelemetryEnabled = clientTelemetryEnabled; } + public boolean isTelemetryServiceAvailable() { + return isTelemetryServiceAvailable; + } + + public void setTelemetryServiceAvailable(boolean isTelemetryServiceAvailable) { + this.isTelemetryServiceAvailable = isTelemetryServiceAvailable; + } + public int getArrayBindStageThreshold() { return arrayBindStageThreshold; } diff --git a/src/main/java/net/snowflake/client/core/SFSession.java b/src/main/java/net/snowflake/client/core/SFSession.java index eafe2263d..a4e89c435 100644 --- a/src/main/java/net/snowflake/client/core/SFSession.java +++ b/src/main/java/net/snowflake/client/core/SFSession.java @@ -63,6 +63,8 @@ public class SFSession extends SFBaseSession { private static final ObjectMapper OBJECT_MAPPER = ObjectMapperFactory.getObjectMapper(); private static final String SF_PATH_SESSION_HEARTBEAT = "/session/heartbeat"; private static final String SF_PATH_QUERY_MONITOR = "/monitoring/queries/"; + public static final String TELEMETRY_SERVICE_AVAILABLE = "TELEMETRY_SERVICE_AVAILABLE"; + // temporarily have this variable to avoid hardcode. // Need to be removed when a better way to organize session parameter is introduced. private static final String CLIENT_STORE_TEMPORARY_CREDENTIAL = @@ -759,6 +761,20 @@ public synchronized void open() throws SFException, SnowflakeSQLException { HttpUtil.setConnectionTimeout(loginInput.getConnectionTimeoutInMillis()); HttpUtil.setSocketTimeout(loginInput.getSocketTimeoutInMillis()); + String telemetryValue = + loginInput + .getSessionParameters() + .getOrDefault(TELEMETRY_SERVICE_AVAILABLE, Boolean.TRUE) + .toString(); + + boolean isTelemetryServiceAvailable = Boolean.parseBoolean(telemetryValue); + + if (!isTelemetryServiceAvailable) { + logger.debug("Telemetry Service is not available"); + } + + setTelemetryServiceAvailable(isTelemetryServiceAvailable); + runDiagnosticsIfEnabled(); SFLoginOutput loginOutput = diff --git a/src/main/java/net/snowflake/client/jdbc/telemetry/TelemetryClient.java b/src/main/java/net/snowflake/client/jdbc/telemetry/TelemetryClient.java index 13b2fb9d2..61d13deef 100644 --- a/src/main/java/net/snowflake/client/jdbc/telemetry/TelemetryClient.java +++ b/src/main/java/net/snowflake/client/jdbc/telemetry/TelemetryClient.java @@ -67,6 +67,7 @@ public class TelemetryClient implements Telemetry { private TelemetryClient(SFSession session, int flushSize) { this.session = session; this.serverUrl = session.getUrl(); + this.isTelemetryServiceAvailable = session.isTelemetryServiceAvailable(); this.httpClient = null; if (this.serverUrl.endsWith("/")) { diff --git a/src/test/java/net/snowflake/client/core/SFLoginInputTest.java b/src/test/java/net/snowflake/client/core/SFLoginInputTest.java index b34eebc02..1e0923442 100644 --- a/src/test/java/net/snowflake/client/core/SFLoginInputTest.java +++ b/src/test/java/net/snowflake/client/core/SFLoginInputTest.java @@ -1,8 +1,14 @@ package net.snowflake.client.core; +import static net.snowflake.client.core.SFSession.TELEMETRY_SERVICE_AVAILABLE; import static org.junit.jupiter.api.Assertions.assertEquals; +import java.util.HashMap; +import java.util.Map; import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.NullSource; +import org.junit.jupiter.params.provider.ValueSource; public class SFLoginInputTest { @@ -19,4 +25,15 @@ public void testGetHostFromServerUrlWithProtocolShouldNotThrow() throws SFExcept sfLoginInput.setServerUrl("https://host.com"); assertEquals("host.com", sfLoginInput.getHostFromServerUrl()); } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + @NullSource + public void testGetTelemetryServiceAvailableShouldReturnCorrectValue(Boolean value) { + SFLoginInput sfLoginInput = new SFLoginInput(); + Map sessionParameters = new HashMap<>(); + sessionParameters.put(TELEMETRY_SERVICE_AVAILABLE, value); + sfLoginInput.setSessionParameters(sessionParameters); + assertEquals(value, sfLoginInput.getSessionParameters().get(TELEMETRY_SERVICE_AVAILABLE)); + } }