implemented close() in DatabricksMetrics

Author: bhu1rdb

src/main/java/com/databricks/jdbc/core/DatabricksSession.java

@@ -12,7 +12,6 @@
 import com.databricks.jdbc.core.types.CompressionType;
 import com.databricks.jdbc.core.types.ComputeResource;
 import com.databricks.jdbc.driver.IDatabricksConnectionContext;
-import com.databricks.jdbc.telemetry.DatabricksMetrics;
 import com.databricks.sdk.support.ToStringer;
 import com.google.common.annotations.VisibleForTesting;
 import java.util.HashMap;
@@ -149,27 +148,14 @@ public void open() throws DatabricksSQLException {
   @Override
   public void close() throws DatabricksSQLException {
     LoggingUtil.log(LogLevel.DEBUG, "public void close()");
-    System.out.println("close() session");
     // TODO: check for any pending query executions
     synchronized (this) {
       if (isSessionOpen) {
         // TODO: handle closed connections by server
         databricksClient.deleteSession(this, computeResource);
         this.sessionInfo = null;
         this.isSessionOpen = false;
-
-        // Flush out metrics when connection is closed
-        DatabricksMetrics metricsExporter = connectionContext.getMetricsExporter();
-        if (metricsExporter != null) {
-          try {
-            metricsExporter.sendRequest(
-                metricsExporter.getGaugeMetrics(), DatabricksMetrics.MetricsType.GAUGE);
-            metricsExporter.sendRequest(
-                metricsExporter.getCounterMetrics(), DatabricksMetrics.MetricsType.COUNTER);
-          } catch (Exception e) {
-            throw new RuntimeException("Failed to flush out metrics", e);
-          }
-        }
+        this.connectionContext.getMetricsExporter().close();
       }
     }
   }

src/main/java/com/databricks/jdbc/telemetry/DatabricksMetrics.java

@@ -1,7 +1,9 @@
 package com.databricks.jdbc.telemetry;
 
 import com.databricks.jdbc.client.DatabricksHttpException;
+import com.databricks.jdbc.commons.LogLevel;
 import com.databricks.jdbc.commons.util.DefaultHttpClientUtil;
+import com.databricks.jdbc.commons.util.LoggingUtil;
 import com.databricks.jdbc.core.DatabricksSQLException;
 import com.databricks.jdbc.driver.IDatabricksConnectionContext;
 import com.fasterxml.jackson.databind.ObjectMapper;
@@ -18,7 +20,7 @@
 import org.apache.http.impl.client.DefaultHttpClient;
 import org.apache.http.util.EntityUtils;
 
-public class DatabricksMetrics {
+public class DatabricksMetrics implements AutoCloseable {
   private final String URL =
       "https://aa87314c1e33d4c1f91a919f8cf9c4ba-387609431.us-west-2.elb.amazonaws.com:443/api/2.0/oss-sql-driver-telemetry/metrics";
   private final Map<String, Double> gaugeMetrics = new HashMap<>();
@@ -32,14 +34,6 @@ public class DatabricksMetrics {
   private String workspaceId = null;
   private DefaultHttpClient telemetryClient = null;
 
-  public Map<String, Double> getGaugeMetrics() {
-    return gaugeMetrics;
-  }
-
-  public Map<String, Double> getCounterMetrics() {
-    return counterMetrics;
-  }
-
   private void setWorkspaceId(String workspaceId) {
     this.workspaceId = workspaceId;
   }
@@ -169,4 +163,19 @@ public void increment(String name, double value) {
     incCounterMetrics(name + "_" + workspaceId, value);
     if (!hasInitialExportOccurred) initialExport(counterMetrics, MetricsType.COUNTER);
   }
+
+  @Override
+  public void close() {
+    // Flush out metrics when connection is closed
+    if (telemetryClient != null) {
+      try {
+        sendRequest(gaugeMetrics, DatabricksMetrics.MetricsType.GAUGE);
+        sendRequest(counterMetrics, DatabricksMetrics.MetricsType.COUNTER);
+      } catch (Exception e) {
+        LoggingUtil.log(
+            LogLevel.DEBUG,
+            "Failed to export metrics when connection is closed. Error: " + e.getMessage());
+      }
+    }
+  }
 }