Flushing metrics out when connection is closed

Author: bhu1rdb

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

@@ -12,6 +12,7 @@
 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;
@@ -148,13 +149,27 @@ 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);
+          }
+        }
       }
     }
   }

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

@@ -18,12 +18,11 @@
 import org.apache.http.impl.client.DefaultHttpClient;
 import org.apache.http.util.EntityUtils;
 
-// TODO (Bhuvan) Flush out metrics once the driver connection is closed
 public class DatabricksMetrics {
   private final String URL =
       "https://aa87314c1e33d4c1f91a919f8cf9c4ba-387609431.us-west-2.elb.amazonaws.com:443/api/2.0/oss-sql-driver-telemetry/metrics";
-  public final Map<String, Double> gaugeMetrics = new HashMap<>();
-  public final Map<String, Double> counterMetrics = new HashMap<>();
+  private final Map<String, Double> gaugeMetrics = new HashMap<>();
+  private final Map<String, Double> counterMetrics = new HashMap<>();
   private final long intervalDurationForSendingReq =
       TimeUnit.SECONDS.toMillis(10 * 60); // 10 minutes
   private final ObjectMapper objectMapper = new ObjectMapper();
@@ -33,6 +32,14 @@ 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;
   }

src/test/java/com/databricks/jdbc/local/DriverTester.java

@@ -166,6 +166,7 @@ void testIfCreateSessionIsExported() throws Exception {
         "jdbc:databricks://e2-dogfood.staging.cloud.databricks.com:443/default;transportMode=https;ssl=1;AuthMech=3;httpPath=/sql/1.0/warehouses/791ba2a31c7fd70a;";
     Connection con = DriverManager.getConnection(jdbcUrl, "user", "x");
     System.out.println("Connection established......");
+    con.close();
   }
 
   @Test