Merge branch 'master' into joey/aws-sns

DataDog · May 23, 2024 · 82f6cdc · 82f6cdc
2 parents 6aa6c6d + 58edd35
commit 82f6cdc
Show file tree

Hide file tree

Showing 5 changed files with 108 additions and 7 deletions.
diff --git a/...nt/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/jdbc/DBInfo.java b/...nt/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/jdbc/DBInfo.java
@@ -13,6 +13,8 @@ public class DBInfo {
   private final String db;
   private final String host;
   private final Integer port;
+  private final String warehouse;
+  private final String schema;
 
   DBInfo(
       String type,
@@ -23,7 +25,9 @@ public class DBInfo {
       String instance,
       String db,
       String host,
-      Integer port) {
+      Integer port,
+      String warehouse,
+      String schema) {
     this.type = type;
     this.subtype = subtype;
     this.fullPropagationSupport = fullPropagationSupport;
@@ -33,6 +37,8 @@ public class DBInfo {
     this.db = db;
     this.host = host;
     this.port = port;
+    this.warehouse = warehouse;
+    this.schema = schema;
   }
 
   public static class Builder {
@@ -45,6 +51,8 @@ public static class Builder {
     private String user;
     private String instance;
     private String db;
+    private String warehouse;
+    private String schema;
     private String host;
     private Integer port;
 
@@ -59,7 +67,9 @@ public static class Builder {
         String instance,
         String db,
         String host,
-        Integer port) {
+        Integer port,
+        String warehouse,
+        String schema) {
       this.type = type;
       this.subtype = subtype;
       this.fullPropagationSupport = fullPropagationSupport;
@@ -69,6 +79,8 @@ public static class Builder {
       this.db = db;
       this.host = host;
       this.port = port;
+      this.warehouse = warehouse;
+      this.schema = schema;
     }
 
     public Builder type(String type) {
@@ -109,6 +121,16 @@ public Builder db(String db) {
       return this;
     }
 
+    public Builder warehouse(String warehouse) {
+      this.warehouse = warehouse;
+      return this;
+    }
+
+    public Builder schema(String schema) {
+      this.schema = schema;
+      return this;
+    }
+
     public Builder host(String host) {
       this.host = host;
       return this;
@@ -120,7 +142,18 @@ public Builder port(Integer port) {
     }
 
     public DBInfo build() {
-      return new DBInfo(type, subtype, fullPropagationSupport, url, user, instance, db, host, port);
+      return new DBInfo(
+          type,
+          subtype,
+          fullPropagationSupport,
+          url,
+          user,
+          instance,
+          db,
+          host,
+          port,
+          warehouse,
+          schema);
     }
   }
 
@@ -160,8 +193,27 @@ public Integer getPort() {
     return port;
   }
 
+  public String getWarehouse() {
+    return warehouse;
+  }
+
+  public String getSchema() {
+    return schema;
+  }
+
   public Builder toBuilder() {
-    return new Builder(type, subtype, fullPropagationSupport, url, user, instance, db, host, port);
+    return new Builder(
+        type,
+        subtype,
+        fullPropagationSupport,
+        url,
+        user,
+        instance,
+        db,
+        host,
+        port,
+        warehouse,
+        schema);
   }
 
   @Override
@@ -177,11 +229,24 @@ public boolean equals(Object o) {
         && Objects.equals(instance, dbInfo.instance)
         && Objects.equals(db, dbInfo.db)
         && Objects.equals(host, dbInfo.host)
-        && Objects.equals(port, dbInfo.port);
+        && Objects.equals(port, dbInfo.port)
+        && Objects.equals(warehouse, dbInfo.warehouse)
+        && Objects.equals(schema, dbInfo.schema);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(type, subtype, fullPropagationSupport, url, user, instance, db, host, port);
+    return Objects.hash(
+        type,
+        subtype,
+        fullPropagationSupport,
+        url,
+        user,
+        instance,
+        db,
+        host,
+        port,
+        warehouse,
+        schema);
   }
 }
diff --git a/...p/src/main/java/datadog/trace/bootstrap/instrumentation/jdbc/JDBCConnectionUrlParser.java b/...p/src/main/java/datadog/trace/bootstrap/instrumentation/jdbc/JDBCConnectionUrlParser.java
@@ -770,6 +770,17 @@ DBInfo.Builder doParse(String jdbcUrl, DBInfo.Builder builder) {
       }
       return builder;
     }
+  },
+
+  SNOWFLAKE("snowflake") {
+    @Override
+    DBInfo.Builder doParse(String jdbcUrl, DBInfo.Builder builder) {
+      String url = jdbcUrl;
+      if (url.startsWith("jdbc:")) {
+        url = url.substring(5);
+      }
+      return GENERIC_URL_LIKE.doParse(url, builder);
+    }
   };
 
   private static final Map<String, JDBCConnectionUrlParser> typeParsers = new HashMap<>();
@@ -876,7 +887,15 @@ private static void populateStandardProperties(
       if (props.containsKey("databaseName")) {
         builder.db((String) props.get("databaseName"));
       }
-
+      if (props.containsKey("db")) {
+        builder.db((String) props.get("db"));
+      }
+      if (props.containsKey("warehouse")) {
+        builder.warehouse((String) props.get("warehouse"));
+      }
+      if (props.containsKey("schema")) {
+        builder.schema((String) props.get("schema"));
+      }
       if (props.containsKey("servername")) {
         builder.host((String) props.get("servername"));
       }

diff --git a/.../instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/JDBCDecorator.java b/.../instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/JDBCDecorator.java
@@ -1,6 +1,8 @@
 package datadog.trace.instrumentation.jdbc;
 
 import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_OPERATION;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_SCHEMA;
+import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_WAREHOUSE;
 
 import datadog.trace.api.Config;
 import datadog.trace.api.DDSpanId;
@@ -119,9 +121,18 @@ protected String dbHostname(final DBInfo info) {
     return info.getHost();
   }
 
+  private void setTagIfPresent(final AgentSpan span, final String key, final String value) {
+    if (value != null && !value.isEmpty()) {
+      span.setTag(key, value);
+    }
+  }
+
   public AgentSpan onConnection(final AgentSpan span, DBInfo dbInfo) {
     if (dbInfo != null) {
       processDatabaseType(span, dbInfo.getType());
+
+      setTagIfPresent(span, DB_WAREHOUSE, dbInfo.getWarehouse());
+      setTagIfPresent(span, DB_SCHEMA, dbInfo.getSchema());
     }
     return super.onConnection(span, dbInfo);
   }

diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/JDBCConnectionUrlParserTest.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/JDBCConnectionUrlParserTest.groovy
@@ -51,6 +51,9 @@ class JDBCConnectionUrlParserTest extends AgentTestRunner {
 
     where:
     url                                                                                                                                                                                                                                         | props    | type         | subtype       | user          | host                                                                | port  | instance                           | db
+    // snowflake
+    "jdbc:snowflake://sza96462.us-east-1.snowflakecomputing.com:443/?db=DATA_OBSERVABILITY_SANDBOX&user=user"                                                                                                                                   | null     | "snowflake"  | null          | "user"        | "sza96462.us-east-1.snowflakecomputing.com"                         | 443   | null                               | "data_observability_sandbox"
+
     // https://jdbc.postgresql.org/documentation/94/connect.html
     "jdbc:postgresql:///"                                                                                                                                                                                                                       | null     | "postgresql" | null          | null          | "localhost"                                                         | 5432  | null                               | null
     "jdbc:postgresql:///"                                                                                                                                                                                                                       | stdProps | "postgresql" | null          | "stdUserName" | "stdServerName"                                                     | 9999  | null                               | "stdDatabaseName"

diff --git a/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Tags.java b/internal-api/src/main/java/datadog/trace/bootstrap/instrumentation/api/Tags.java
@@ -42,6 +42,9 @@ public class Tags {
   public static final String DB_USER = "db.user";
   public static final String DB_OPERATION = "db.operation";
   public static final String DB_STATEMENT = "db.statement";
+  public static final String DB_WAREHOUSE = "db.warehouse";
+  public static final String DB_HOST = "db.host";
+  public static final String DB_SCHEMA = "db.schema";
   public static final String MESSAGE_BUS_DESTINATION = "message_bus.destination";
 
   public static final String TEST_MODULE = "test.module";