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 5201175c4..b302edb1c 100644 --- a/src/main/java/net/snowflake/client/jdbc/telemetry/TelemetryClient.java +++ b/src/main/java/net/snowflake/client/jdbc/telemetry/TelemetryClient.java @@ -184,11 +184,26 @@ public static Telemetry createTelemetry(SFSession session, int flushSize) { return new TelemetryClient(session, flushSize); } + /** + * Initialize the sessionless telemetry connector using KEYPAIR_JWT as the default auth type + * + * @param httpClient client object used to communicate with other machine + * @param serverUrl server url + * @return a telemetry connector + */ + public static Telemetry createSessionlessTelemetry( + CloseableHttpClient httpClient, String serverUrl) { + // By default, use KEYPAIR_JWT as the auth type + return createSessionlessTelemetry( + httpClient, serverUrl, "KEYPAIR_JWT", DEFAULT_FORCE_FLUSH_SIZE); + } + /** * Initialize the sessionless telemetry connector * * @param httpClient client object used to communicate with other machine * @param serverUrl server url + * @param authType authorization type for sessionless telemetry * @return a telemetry connector */ public static Telemetry createSessionlessTelemetry( diff --git a/src/test/java/net/snowflake/client/jdbc/telemetry/TelemetryIT.java b/src/test/java/net/snowflake/client/jdbc/telemetry/TelemetryIT.java index 63fce6782..c9b9b6272 100644 --- a/src/test/java/net/snowflake/client/jdbc/telemetry/TelemetryIT.java +++ b/src/test/java/net/snowflake/client/jdbc/telemetry/TelemetryIT.java @@ -47,6 +47,13 @@ public void testTelemetry() throws Exception { testTelemetryInternal(telemetry); } + @Ignore + @Test + @ConditionalIgnoreRule.ConditionalIgnore(condition = RunningOnGithubAction.class) + public void testSessionlessTelemetry() throws Exception, SFException { + testTelemetryInternal(createSessionlessTelemetry()); + } + @Ignore @Test @ConditionalIgnoreRule.ConditionalIgnore(condition = RunningOnGithubAction.class) @@ -197,6 +204,26 @@ public void testClosedOAuthSessionlessTelemetry() throws Exception, SFException Assert.assertFalse(telemetry.sendBatchAsync().get()); } + // Helper function to create a sessionless telemetry + // using createSessionlessTelemetry(CloseableHttpClient httpClient, String serverUrl) + private TelemetryClient createSessionlessTelemetry() + throws SFException, SQLException, IOException { + setUpPublicKey(); + String privateKeyLocation = getFullPathFileInResource("rsa_key.p8"); + Map parameters = getConnectionParameters(); + String jwtToken = + SessionUtil.generateJWTToken( + null, privateKeyLocation, null, parameters.get("account"), parameters.get("user")); + + CloseableHttpClient httpClient = HttpUtil.buildHttpClient(null, null, false); + TelemetryClient telemetry = + (TelemetryClient) + TelemetryClient.createSessionlessTelemetry( + httpClient, String.format("%s:%s", parameters.get("host"), parameters.get("port"))); + telemetry.refreshToken(jwtToken); + return telemetry; + } + // Helper function to create a sessionless telemetry using keypair JWT private TelemetryClient createJWTSessionlessTelemetry() throws SFException, SQLException, IOException {