SNOW-1016363 Add rich results metadata support for Coldbrew/Snowsight…

… integration (#1840)
snowflakedb · Jul 23, 2024 · 14b62bd · 14b62bd
1 parent 6392f61
commit 14b62bd
Show file tree

Hide file tree

Showing 6 changed files with 872 additions and 248 deletions.
diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeColumnMetadata.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeColumnMetadata.java
@@ -4,8 +4,16 @@
 
 package net.snowflake.client.jdbc;
 
+import static net.snowflake.client.jdbc.SnowflakeUtil.getFieldMetadata;
+import static net.snowflake.client.jdbc.SnowflakeUtil.getSnowflakeType;
+import static net.snowflake.client.jdbc.SnowflakeUtil.isVectorType;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.google.common.base.Strings;
 import java.io.Serializable;
+import java.sql.Types;
 import java.util.List;
+import net.snowflake.client.core.SFBaseSession;
 import net.snowflake.client.core.SnowflakeJdbcInternalApi;
 
 /**
@@ -99,6 +107,57 @@ public SnowflakeColumnMetadata(
     this.isAutoIncrement = isAutoIncrement;
   }
 
+  @SnowflakeJdbcInternalApi
+  public SnowflakeColumnMetadata(
+      JsonNode colNode, boolean jdbcTreatDecimalAsInt, SFBaseSession session)
+      throws SnowflakeSQLLoggedException {
+    this.name = colNode.path("name").asText();
+    this.nullable = colNode.path("nullable").asBoolean();
+    this.precision = colNode.path("precision").asInt();
+    this.scale = colNode.path("scale").asInt();
+    this.length = colNode.path("length").asInt();
+    int dimension =
+        colNode
+            .path("dimension")
+            .asInt(); // vector dimension when checking columns via connection.getMetadata
+    int vectorDimension =
+        colNode
+            .path("vectorDimension")
+            .asInt(); // dimension when checking columns via resultSet.getMetadata
+    this.dimension = dimension > 0 ? dimension : vectorDimension;
+    this.fixed = colNode.path("fixed").asBoolean();
+    JsonNode udtOutputType = colNode.path("outputType");
+    JsonNode extColTypeNameNode = colNode.path("extTypeName");
+    String extColTypeName = null;
+    if (!extColTypeNameNode.isMissingNode()
+        && !Strings.isNullOrEmpty(extColTypeNameNode.asText())) {
+      extColTypeName = extColTypeNameNode.asText();
+    }
+    String internalColTypeName = colNode.path("type").asText();
+    List<FieldMetadata> fieldsMetadata =
+        getFieldMetadata(jdbcTreatDecimalAsInt, internalColTypeName, colNode);
+
+    int fixedColType = jdbcTreatDecimalAsInt && scale == 0 ? Types.BIGINT : Types.DECIMAL;
+    ColumnTypeInfo columnTypeInfo =
+        getSnowflakeType(
+            internalColTypeName,
+            extColTypeName,
+            udtOutputType,
+            session,
+            fixedColType,
+            !fieldsMetadata.isEmpty(),
+            isVectorType(internalColTypeName));
+
+    this.typeName = columnTypeInfo.getExtColTypeName();
+    this.type = columnTypeInfo.getColumnType();
+    this.base = columnTypeInfo.getSnowflakeType();
+    this.fields = fieldsMetadata;
+    this.columnSrcDatabase = colNode.path("database").asText();
+    this.columnSrcSchema = colNode.path("schema").asText();
+    this.columnSrcTable = colNode.path("table").asText();
+    this.isAutoIncrement = colNode.path("isAutoIncrement").asBoolean();
+  }
+
   public String getName() {
     return name;
   }

diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeDatabaseMetaData.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeDatabaseMetaData.java
@@ -1825,8 +1825,7 @@ public boolean next() throws SQLException {
             logger.debug("Data type string: {}", dataTypeStr);
 
             SnowflakeColumnMetadata columnMetadata =
-                SnowflakeUtil.extractColumnMetadata(
-                    jsonNode, session.isJdbcTreatDecimalAsInt(), session);
+                new SnowflakeColumnMetadata(jsonNode, session.isJdbcTreatDecimalAsInt(), session);
 
             logger.debug("Nullable: {}", columnMetadata.isNullable());