diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeColumnMetadata.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeColumnMetadata.java index 9f1cd272e..4525c2efb 100644 --- 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 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 index acfb3e4f7..88ee255e1 100644 --- 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()); diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeResultSetSerializableV1.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeResultSetSerializableV1.java index add20978f..535019977 100644 --- a/src/main/java/net/snowflake/client/jdbc/SnowflakeResultSetSerializableV1.java +++ b/src/main/java/net/snowflake/client/jdbc/SnowflakeResultSetSerializableV1.java @@ -12,6 +12,7 @@ import static net.snowflake.client.core.SessionUtil.CLIENT_RESULT_CHUNK_SIZE; import static net.snowflake.client.core.SessionUtil.DEFAULT_CLIENT_MEMORY_LIMIT; import static net.snowflake.client.core.SessionUtil.DEFAULT_CLIENT_PREFETCH_THREADS; +import static net.snowflake.client.jdbc.SnowflakeChunkDownloader.NoOpChunkDownloader; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; @@ -46,6 +47,7 @@ import net.snowflake.client.core.SFSession; import net.snowflake.client.core.SFStatementType; import net.snowflake.client.core.SessionUtil; +import net.snowflake.client.core.SnowflakeJdbcInternalApi; import net.snowflake.client.jdbc.telemetry.NoOpTelemetryClient; import net.snowflake.client.jdbc.telemetry.Telemetry; import net.snowflake.client.log.ArgSupplier; @@ -272,6 +274,199 @@ private SnowflakeResultSetSerializableV1(SnowflakeResultSetSerializableV1 toCopy this.resultStreamProvider = toCopy.resultStreamProvider; } + /** + * @param rootNode result JSON node received from GS + * @param sfSession the Snowflake session + * @param sfStatement the Snowflake statement + * @param resultStreamProvider a ResultStreamProvider for computing a custom data source for + * result-file streams + * @throws SnowflakeSQLException if failed to parse the result JSON node + */ + protected SnowflakeResultSetSerializableV1( + JsonNode rootNode, + SFBaseSession sfSession, + SFBaseStatement sfStatement, + ResultStreamProvider resultStreamProvider, + boolean disableChunksPrefetch) + throws SnowflakeSQLException { + SnowflakeUtil.checkErrorAndThrowException(rootNode); + + // get the query id + this.queryId = rootNode.path("data").path("queryId").asText(); + + JsonNode databaseNode = rootNode.path("data").path("finalDatabaseName"); + this.finalDatabaseName = + databaseNode.isNull() + ? (sfSession != null ? sfSession.getDatabase() : null) + : databaseNode.asText(); + + JsonNode schemaNode = rootNode.path("data").path("finalSchemaName"); + this.finalSchemaName = + schemaNode.isNull() + ? (sfSession != null ? sfSession.getSchema() : null) + : schemaNode.asText(); + + JsonNode roleNode = rootNode.path("data").path("finalRoleName"); + this.finalRoleName = + roleNode.isNull() ? (sfSession != null ? sfSession.getRole() : null) : roleNode.asText(); + + JsonNode warehouseNode = rootNode.path("data").path("finalWarehouseName"); + this.finalWarehouseName = + warehouseNode.isNull() + ? (sfSession != null ? sfSession.getWarehouse() : null) + : warehouseNode.asText(); + + this.statementType = + SFStatementType.lookUpTypeById(rootNode.path("data").path("statementTypeId").asLong()); + + this.totalRowCountTruncated = rootNode.path("data").path("totalTruncated").asBoolean(); + + this.possibleSession = Optional.ofNullable(sfSession); + + logger.debug("Query id: {}", this.queryId); + + Optional queryResultFormat = + QueryResultFormat.lookupByName(rootNode.path("data").path("queryResultFormat").asText()); + this.queryResultFormat = queryResultFormat.orElse(QueryResultFormat.JSON); + + // extract query context and save it in current session + JsonNode queryContextNode = rootNode.path("data").path("queryContext"); + String queryContext = queryContextNode.isNull() ? null : queryContextNode.toString(); + + if (!sfSession.isAsyncSession()) { + sfSession.setQueryContext(queryContext); + } + + // extract parameters + this.parameters = SessionUtil.getCommonParams(rootNode.path("data").path("parameters")); + if (this.parameters.isEmpty()) { + this.parameters = new HashMap<>(sfSession.getCommonParameters()); + this.setStatemementLevelParameters(sfStatement.getStatementParameters()); + } + + // initialize column metadata + this.columnCount = rootNode.path("data").path("rowtype").size(); + + for (int i = 0; i < this.columnCount; i++) { + JsonNode colNode = rootNode.path("data").path("rowtype").path(i); + + SnowflakeColumnMetadata columnMetadata = + new SnowflakeColumnMetadata(colNode, sfSession.isJdbcTreatDecimalAsInt(), sfSession); + + this.resultColumnMetadata.add(columnMetadata); + + logger.debug("Get column metadata: {}", (ArgSupplier) columnMetadata::toString); + } + + this.resultStreamProvider = resultStreamProvider; + + // process the content of first chunk. + if (this.queryResultFormat == QueryResultFormat.ARROW) { + this.firstChunkStringData = rootNode.path("data").path("rowsetBase64").asText(); + this.rootAllocator = new RootAllocator(Long.MAX_VALUE); + // Set first chunk row count from firstChunkStringData + this.setFirstChunkRowCountForArrow(); + } else { + this.firstChunkRowset = rootNode.path("data").path("rowset"); + + if (this.firstChunkRowset == null || this.firstChunkRowset.isMissingNode()) { + this.firstChunkRowCount = 0; + this.firstChunkStringData = null; + this.firstChunkByteData = new byte[0]; + } else { + this.firstChunkRowCount = this.firstChunkRowset.size(); + this.firstChunkStringData = this.firstChunkRowset.toString(); + } + } + logger.debug("First chunk row count: {}", this.firstChunkRowCount); + + // parse file chunks + this.parseChunkFiles(rootNode, sfStatement); + + // result version + JsonNode versionNode = rootNode.path("data").path("version"); + + if (!versionNode.isMissingNode()) { + this.resultVersion = versionNode.longValue(); + } + + // number of binds + JsonNode numberOfBindsNode = rootNode.path("data").path("numberOfBinds"); + + if (!numberOfBindsNode.isMissingNode()) { + this.numberOfBinds = numberOfBindsNode.intValue(); + } + + JsonNode arrayBindSupported = rootNode.path("data").path("arrayBindSupported"); + this.arrayBindSupported = !arrayBindSupported.isMissingNode() && arrayBindSupported.asBoolean(); + + // time result sent by GS (epoch time in millis) + JsonNode sendResultTimeNode = rootNode.path("data").path("sendResultTime"); + if (!sendResultTimeNode.isMissingNode()) { + this.sendResultTime = sendResultTimeNode.longValue(); + } + + logger.debug("Result version: {}", this.resultVersion); + + // Bind parameter metadata + JsonNode bindData = rootNode.path("data").path("metaDataOfBinds"); + if (!bindData.isMissingNode()) { + List returnVal = new ArrayList<>(); + for (JsonNode child : bindData) { + int precision = child.path("precision").asInt(); + boolean nullable = child.path("nullable").asBoolean(); + int scale = child.path("scale").asInt(); + int byteLength = child.path("byteLength").asInt(); + int length = child.path("length").asInt(); + String name = child.path("name").asText(); + String type = child.path("type").asText(); + MetaDataOfBinds param = + new MetaDataOfBinds(precision, nullable, scale, byteLength, length, name, type); + returnVal.add(param); + } + this.metaDataOfBinds = returnVal; + } + + // setup fields from sessions. + this.ocspMode = sfSession.getOCSPMode(); + this.httpClientKey = sfSession.getHttpClientKey(); + this.snowflakeConnectionString = sfSession.getSnowflakeConnectionString(); + this.networkTimeoutInMilli = sfSession.getNetworkTimeoutInMilli(); + this.authTimeout = 0; + this.maxHttpRetries = sfSession.getMaxHttpRetries(); + this.isResultColumnCaseInsensitive = sfSession.isResultColumnCaseInsensitive(); + this.treatNTZAsUTC = sfSession.getTreatNTZAsUTC(); + this.formatDateWithTimezone = sfSession.getFormatDateWithTimezone(); + this.useSessionTimezone = sfSession.getUseSessionTimezone(); + + // setup transient fields from parameter + this.setupFieldsFromParameters(); + + if (disableChunksPrefetch) { + this.chunkDownloader = new NoOpChunkDownloader(); + } else { + this.chunkDownloader = + (this.chunkFileCount > 0) + // The chunk downloader will start prefetching + // first few chunk files in background thread(s) + ? new SnowflakeChunkDownloader(this) + : new NoOpChunkDownloader(); + } + + // Setup ResultSet metadata + this.resultSetMetaData = + new SFResultSetMetaData( + this.getResultColumnMetadata(), + this.queryId, + sfSession, + this.isResultColumnCaseInsensitive, + this.timestampNTZFormatter, + this.timestampLTZFormatter, + this.timestampTZFormatter, + this.dateFormatter, + this.timeFormatter); + } + public void setRootAllocator(RootAllocator rootAllocator) { this.rootAllocator = rootAllocator; } @@ -544,190 +739,22 @@ public static SnowflakeResultSetSerializableV1 create( SFBaseStatement sfStatement, ResultStreamProvider resultStreamProvider) throws SnowflakeSQLException { - SnowflakeResultSetSerializableV1 resultSetSerializable = new SnowflakeResultSetSerializableV1(); logger.trace("Entering create()", false); + return new SnowflakeResultSetSerializableV1( + rootNode, sfSession, sfStatement, resultStreamProvider, false); + } - SnowflakeUtil.checkErrorAndThrowException(rootNode); - - // get the query id - resultSetSerializable.queryId = rootNode.path("data").path("queryId").asText(); - - JsonNode databaseNode = rootNode.path("data").path("finalDatabaseName"); - resultSetSerializable.finalDatabaseName = - databaseNode.isNull() - ? (sfSession != null ? sfSession.getDatabase() : null) - : databaseNode.asText(); - - JsonNode schemaNode = rootNode.path("data").path("finalSchemaName"); - resultSetSerializable.finalSchemaName = - schemaNode.isNull() - ? (sfSession != null ? sfSession.getSchema() : null) - : schemaNode.asText(); - - JsonNode roleNode = rootNode.path("data").path("finalRoleName"); - resultSetSerializable.finalRoleName = - roleNode.isNull() ? (sfSession != null ? sfSession.getRole() : null) : roleNode.asText(); - - JsonNode warehouseNode = rootNode.path("data").path("finalWarehouseName"); - resultSetSerializable.finalWarehouseName = - warehouseNode.isNull() - ? (sfSession != null ? sfSession.getWarehouse() : null) - : warehouseNode.asText(); - - resultSetSerializable.statementType = - SFStatementType.lookUpTypeById(rootNode.path("data").path("statementTypeId").asLong()); - - resultSetSerializable.totalRowCountTruncated = - rootNode.path("data").path("totalTruncated").asBoolean(); - - resultSetSerializable.possibleSession = Optional.ofNullable(sfSession); - - logger.debug("Query id: {}", resultSetSerializable.queryId); - - Optional queryResultFormat = - QueryResultFormat.lookupByName(rootNode.path("data").path("queryResultFormat").asText()); - resultSetSerializable.queryResultFormat = queryResultFormat.orElse(QueryResultFormat.JSON); - - // extract query context and save it in current session - JsonNode queryContextNode = rootNode.path("data").path("queryContext"); - String queryContext = queryContextNode.isNull() ? null : queryContextNode.toString(); - - if (!sfSession.isAsyncSession()) { - sfSession.setQueryContext(queryContext); - } - - // extract parameters - resultSetSerializable.parameters = - SessionUtil.getCommonParams(rootNode.path("data").path("parameters")); - if (resultSetSerializable.parameters.isEmpty()) { - resultSetSerializable.parameters = new HashMap<>(sfSession.getCommonParameters()); - resultSetSerializable.setStatemementLevelParameters(sfStatement.getStatementParameters()); - } - - // initialize column metadata - resultSetSerializable.columnCount = rootNode.path("data").path("rowtype").size(); - - for (int i = 0; i < resultSetSerializable.columnCount; i++) { - JsonNode colNode = rootNode.path("data").path("rowtype").path(i); - - SnowflakeColumnMetadata columnMetadata = - SnowflakeUtil.extractColumnMetadata( - colNode, sfSession.isJdbcTreatDecimalAsInt(), sfSession); - - resultSetSerializable.resultColumnMetadata.add(columnMetadata); - - logger.debug("Get column metadata: {}", (ArgSupplier) () -> columnMetadata.toString()); - } - - resultSetSerializable.resultStreamProvider = resultStreamProvider; - - // process the content of first chunk. - if (resultSetSerializable.queryResultFormat == QueryResultFormat.ARROW) { - resultSetSerializable.firstChunkStringData = - rootNode.path("data").path("rowsetBase64").asText(); - resultSetSerializable.rootAllocator = new RootAllocator(Long.MAX_VALUE); - // Set first chunk row count from firstChunkStringData - resultSetSerializable.setFirstChunkRowCountForArrow(); - } else { - resultSetSerializable.firstChunkRowset = rootNode.path("data").path("rowset"); - - if (resultSetSerializable.firstChunkRowset == null - || resultSetSerializable.firstChunkRowset.isMissingNode()) { - resultSetSerializable.firstChunkRowCount = 0; - resultSetSerializable.firstChunkStringData = null; - resultSetSerializable.firstChunkByteData = new byte[0]; - } else { - resultSetSerializable.firstChunkRowCount = resultSetSerializable.firstChunkRowset.size(); - resultSetSerializable.firstChunkStringData = - resultSetSerializable.firstChunkRowset.toString(); - } - } - logger.debug("First chunk row count: {}", resultSetSerializable.firstChunkRowCount); - - // parse file chunks - resultSetSerializable.parseChunkFiles(rootNode, sfStatement); - - // result version - JsonNode versionNode = rootNode.path("data").path("version"); - - if (!versionNode.isMissingNode()) { - resultSetSerializable.resultVersion = versionNode.longValue(); - } - - // number of binds - JsonNode numberOfBindsNode = rootNode.path("data").path("numberOfBinds"); - - if (!numberOfBindsNode.isMissingNode()) { - resultSetSerializable.numberOfBinds = numberOfBindsNode.intValue(); - } - - JsonNode arrayBindSupported = rootNode.path("data").path("arrayBindSupported"); - resultSetSerializable.arrayBindSupported = - !arrayBindSupported.isMissingNode() && arrayBindSupported.asBoolean(); - - // time result sent by GS (epoch time in millis) - JsonNode sendResultTimeNode = rootNode.path("data").path("sendResultTime"); - if (!sendResultTimeNode.isMissingNode()) { - resultSetSerializable.sendResultTime = sendResultTimeNode.longValue(); - } - - logger.debug("Result version: {}", resultSetSerializable.resultVersion); - - // Bind parameter metadata - JsonNode bindData = rootNode.path("data").path("metaDataOfBinds"); - if (!bindData.isMissingNode()) { - List returnVal = new ArrayList<>(); - for (JsonNode child : bindData) { - int precision = child.path("precision").asInt(); - boolean nullable = child.path("nullable").asBoolean(); - int scale = child.path("scale").asInt(); - int byteLength = child.path("byteLength").asInt(); - int length = child.path("length").asInt(); - String name = child.path("name").asText(); - String type = child.path("type").asText(); - MetaDataOfBinds param = - new MetaDataOfBinds(precision, nullable, scale, byteLength, length, name, type); - returnVal.add(param); - } - resultSetSerializable.metaDataOfBinds = returnVal; - } - - // setup fields from sessions. - resultSetSerializable.ocspMode = sfSession.getOCSPMode(); - resultSetSerializable.httpClientKey = sfSession.getHttpClientKey(); - resultSetSerializable.snowflakeConnectionString = sfSession.getSnowflakeConnectionString(); - resultSetSerializable.networkTimeoutInMilli = sfSession.getNetworkTimeoutInMilli(); - resultSetSerializable.authTimeout = 0; - resultSetSerializable.maxHttpRetries = sfSession.getMaxHttpRetries(); - resultSetSerializable.isResultColumnCaseInsensitive = sfSession.isResultColumnCaseInsensitive(); - resultSetSerializable.treatNTZAsUTC = sfSession.getTreatNTZAsUTC(); - resultSetSerializable.formatDateWithTimezone = sfSession.getFormatDateWithTimezone(); - resultSetSerializable.useSessionTimezone = sfSession.getUseSessionTimezone(); - - // setup transient fields from parameter - resultSetSerializable.setupFieldsFromParameters(); - - // The chunk downloader will start prefetching - // first few chunk files in background thread(s) - resultSetSerializable.chunkDownloader = - (resultSetSerializable.chunkFileCount > 0) - ? new SnowflakeChunkDownloader(resultSetSerializable) - : new SnowflakeChunkDownloader.NoOpChunkDownloader(); - - // Setup ResultSet metadata - resultSetSerializable.resultSetMetaData = - new SFResultSetMetaData( - resultSetSerializable.getResultColumnMetadata(), - resultSetSerializable.queryId, - sfSession, - resultSetSerializable.isResultColumnCaseInsensitive, - resultSetSerializable.timestampNTZFormatter, - resultSetSerializable.timestampLTZFormatter, - resultSetSerializable.timestampTZFormatter, - resultSetSerializable.dateFormatter, - resultSetSerializable.timeFormatter); - - return resultSetSerializable; + /** + * A factory function for internal usage only. It creates SnowflakeResultSetSerializableV1 with + * NoOpChunksDownloader which disables chunks prefetch. + */ + @SnowflakeJdbcInternalApi + public static SnowflakeResultSetSerializableV1 createWithChunksPrefetchDisabled( + JsonNode rootNode, SFBaseSession sfSession, SFBaseStatement sfStatement) + throws SnowflakeSQLException { + logger.trace("Entering create()", false); + return new SnowflakeResultSetSerializableV1( + rootNode, sfSession, sfStatement, new DefaultResultStreamProvider(), true); } /** @@ -975,9 +1002,7 @@ private void setupTransientFields() throws SQLException { // Allocate chunk downloader if necessary chunkDownloader = - (this.chunkFileCount > 0) - ? new SnowflakeChunkDownloader(this) - : new SnowflakeChunkDownloader.NoOpChunkDownloader(); + (this.chunkFileCount > 0) ? new SnowflakeChunkDownloader(this) : new NoOpChunkDownloader(); this.possibleSession = Optional.empty(); // we don't have session object during deserializing } diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeRichResultSetSerializableV1.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeRichResultSetSerializableV1.java new file mode 100644 index 000000000..194748317 --- /dev/null +++ b/src/main/java/net/snowflake/client/jdbc/SnowflakeRichResultSetSerializableV1.java @@ -0,0 +1,216 @@ +package net.snowflake.client.jdbc; + +import com.fasterxml.jackson.databind.JsonNode; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import net.snowflake.client.core.QueryResultFormat; +import net.snowflake.client.core.SFBaseSession; +import net.snowflake.client.core.SFBaseStatement; +import net.snowflake.client.core.SnowflakeJdbcInternalApi; +import net.snowflake.client.log.ArgSupplier; +import net.snowflake.client.log.SFLogger; +import net.snowflake.client.log.SFLoggerFactory; + +@SnowflakeJdbcInternalApi +public class SnowflakeRichResultSetSerializableV1 extends SnowflakeResultSetSerializableV1 { + + private static final SFLogger logger = + SFLoggerFactory.getLogger(SnowflakeRichResultSetSerializableV1.class); + + private String richResultsFirstChunkStringData; + private int richResultsFirstChunkRowCount; + private int richResultsChunkFileCount; + private int richResultsColumnCount; + private final List richResultsChunkFilesMetadata = new ArrayList<>(); + private byte[] richResultsFirstChunkByteData; + private String richResultsQrmk; + private final Map richResultsChunkHeadersMap = new HashMap<>(); + private final List richResultsColumnMetadata = + new ArrayList<>(); + private QueryResultFormat richResultsQueryResultFormat; + + transient JsonNode richResultFirstChunkRowset = null; + + /** + * A factory function for internal usage only. It creates SnowflakeRichResultSetSerializableV1 + * with NoOpChunksDownloader which disables chunks prefetch. + */ + public static SnowflakeRichResultSetSerializableV1 createWithChunksPrefetchDisabled( + JsonNode rootNode, SFBaseSession sfSession, SFBaseStatement sfStatement) + throws SnowflakeSQLException { + return new SnowflakeRichResultSetSerializableV1( + rootNode, sfSession, sfStatement, new DefaultResultStreamProvider(), true); + } + + private SnowflakeRichResultSetSerializableV1( + JsonNode rootNode, + SFBaseSession sfSession, + SFBaseStatement sfStatement, + ResultStreamProvider resultStreamProvider, + boolean disableChunksPrefetch) + throws SnowflakeSQLException { + super(rootNode, sfSession, sfStatement, resultStreamProvider, disableChunksPrefetch); + if (!rootNode.at("/richResult").isMissingNode()) { + JsonNode richResultsNode = rootNode.path("richResult"); + Optional queryResultFormat = + QueryResultFormat.lookupByName(richResultsNode.path("queryResultFormat").asText()); + this.richResultsQueryResultFormat = queryResultFormat.orElse(QueryResultFormat.JSON); + initializeColumnMetadata(richResultsNode, sfSession); + initializeFirstChunkData(richResultsNode, this.queryResultFormat); + initializeChunkFiles(richResultsNode); + } else { + logger.debug("Unable to initialize rich results metadata, no \"richResult\" node"); + } + } + + private void initializeColumnMetadata(JsonNode richResultsNode, SFBaseSession sfSession) + throws SnowflakeSQLException { + this.richResultsColumnCount = richResultsNode.path("rowtype").size(); + for (int i = 0; i < this.richResultsColumnCount; i++) { + JsonNode colNode = richResultsNode.path("rowtype").path(i); + + SnowflakeRichResultsColumnMetadata columnMetadata = + new SnowflakeRichResultsColumnMetadata( + colNode, sfSession.isJdbcTreatDecimalAsInt(), sfSession); + + this.richResultsColumnMetadata.add(columnMetadata); + logger.debug("Get column metadata: {}", (ArgSupplier) columnMetadata::toString); + } + } + + private void initializeFirstChunkData( + JsonNode richResultsNode, QueryResultFormat queryResultFormat) { + if (queryResultFormat == QueryResultFormat.ARROW) { + this.richResultsFirstChunkStringData = richResultsNode.path("rowsetBase64").asText(); + } else { + this.richResultFirstChunkRowset = richResultsNode.path("rowset"); + + if (this.richResultFirstChunkRowset == null + || this.richResultFirstChunkRowset.isMissingNode()) { + this.richResultsFirstChunkRowCount = 0; + this.richResultsFirstChunkStringData = null; + this.richResultsFirstChunkByteData = new byte[0]; + } else { + this.richResultsFirstChunkRowCount = this.richResultFirstChunkRowset.size(); + this.richResultsFirstChunkStringData = this.richResultFirstChunkRowset.toString(); + } + logger.debug("First rich results chunk row count: {}", this.richResultsFirstChunkRowCount); + } + } + + private void initializeChunkFiles(JsonNode richResultsNode) { + JsonNode chunksNode = richResultsNode.path("chunks"); + if (!chunksNode.isMissingNode()) { + this.richResultsChunkFileCount = chunksNode.size(); + JsonNode qrmkNode = richResultsNode.path("qrmk"); + this.richResultsQrmk = qrmkNode.isMissingNode() ? null : qrmkNode.textValue(); + if (this.richResultsChunkFileCount > 0) { + logger.debug("Number of rich results metadata chunks: {}", this.richResultsChunkFileCount); + initializeChunkHeaders(richResultsNode); + initializeChunkFilesMetadata(chunksNode); + } + } + } + + private void initializeChunkHeaders(JsonNode richResultsNode) { + JsonNode chunkHeaders = richResultsNode.path("chunkHeaders"); + if (chunkHeaders != null && !chunkHeaders.isMissingNode()) { + Iterator> chunkHeadersIter = chunkHeaders.fields(); + while (chunkHeadersIter.hasNext()) { + Map.Entry chunkHeader = chunkHeadersIter.next(); + logger.debug( + "Add header key: {}, value: {}", chunkHeader.getKey(), chunkHeader.getValue().asText()); + this.richResultsChunkHeadersMap.put(chunkHeader.getKey(), chunkHeader.getValue().asText()); + } + } + } + + private void initializeChunkFilesMetadata(JsonNode chunksNode) { + for (int idx = 0; idx < this.richResultsChunkFileCount; idx++) { + JsonNode chunkNode = chunksNode.get(idx); + String url = chunkNode.path("url").asText(); + int rowCount = chunkNode.path("rowCount").asInt(); + int compressedSize = chunkNode.path("compressedSize").asInt(); + int uncompressedSize = chunkNode.path("uncompressedSize").asInt(); + this.richResultsChunkFilesMetadata.add( + new ChunkFileMetadata(url, rowCount, compressedSize, uncompressedSize)); + logger.debug( + "Add rich results metadata chunk, url: {} rowCount: {} " + + "compressedSize: {} uncompressedSize: {}", + url, + rowCount, + compressedSize, + uncompressedSize); + } + } + + public String getRichResultsFirstChunkStringData() { + return richResultsFirstChunkStringData; + } + + public int getRichResultsFirstChunkRowCount() { + return richResultsFirstChunkRowCount; + } + + public int getRichResultsChunkFileCount() { + return richResultsChunkFileCount; + } + + public int getRichResultsColumnCount() { + return richResultsColumnCount; + } + + public List getRichResultsChunkFilesMetadata() { + return richResultsChunkFilesMetadata; + } + + public byte[] getRichResultsFirstChunkByteData() { + return richResultsFirstChunkByteData; + } + + public String getRichResultsQrmk() { + return richResultsQrmk; + } + + public Map getRichResultsChunkHeadersMap() { + return richResultsChunkHeadersMap; + } + + public List getRichResultsColumnMetadata() { + return richResultsColumnMetadata; + } + + public QueryResultFormat getRichResultsQueryResultFormat() { + return richResultsQueryResultFormat; + } + + public JsonNode getRichResultFirstChunkRowset() { + return richResultFirstChunkRowset; + } + + @SnowflakeJdbcInternalApi + public static class SnowflakeRichResultsColumnMetadata extends SnowflakeColumnMetadata { + + private final int columnIndex; + + public SnowflakeRichResultsColumnMetadata( + JsonNode colNode, boolean jdbcTreatDecimalAsInt, SFBaseSession session) + throws SnowflakeSQLLoggedException { + super(colNode, jdbcTreatDecimalAsInt, session); + this.columnIndex = colNode.path("columnIndexing").asInt(); + } + + public int getColumnIndex() { + return columnIndex; + } + + @Override + public String toString() { + return super.toString() + ",columnIndex=" + columnIndex; + } + } +} diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeUtil.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeUtil.java index efe8ffbf8..1485249b3 100644 --- a/src/main/java/net/snowflake/client/jdbc/SnowflakeUtil.java +++ b/src/main/java/net/snowflake/client/jdbc/SnowflakeUtil.java @@ -172,68 +172,12 @@ private static void checkErrorAndThrowExceptionSub( throw new SnowflakeSQLException(queryId, errorMessage, sqlState, errorCode); } + /** This method should only be used internally */ + @Deprecated public static SnowflakeColumnMetadata extractColumnMetadata( JsonNode colNode, boolean jdbcTreatDecimalAsInt, SFBaseSession session) throws SnowflakeSQLException { - String colName = colNode.path("name").asText(); - String internalColTypeName = colNode.path("type").asText(); - boolean nullable = colNode.path("nullable").asBoolean(); - int precision = colNode.path("precision").asInt(); - int scale = colNode.path("scale").asInt(); - int 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 - int finalVectorDimension = dimension > 0 ? dimension : vectorDimension; - boolean 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(); - } - List fieldsMetadata = - getFieldMetadata(jdbcTreatDecimalAsInt, internalColTypeName, colNode); - - int fixedColType = jdbcTreatDecimalAsInt && scale == 0 ? Types.BIGINT : Types.DECIMAL; - ColumnTypeInfo columnTypeInfo = - getSnowflakeType( - internalColTypeName, - extColTypeName, - udtOutputType, - session, - fixedColType, - fieldsMetadata.size() > 0, - isVectorType(internalColTypeName)); - - String colSrcDatabase = colNode.path("database").asText(); - String colSrcSchema = colNode.path("schema").asText(); - String colSrcTable = colNode.path("table").asText(); - - boolean isAutoIncrement = colNode.path("isAutoIncrement").asBoolean(); - - return new SnowflakeColumnMetadata( - colName, - columnTypeInfo.getColumnType(), - nullable, - length, - precision, - scale, - columnTypeInfo.getExtColTypeName(), - fixed, - columnTypeInfo.getSnowflakeType(), - fieldsMetadata, - colSrcDatabase, - colSrcSchema, - colSrcTable, - isAutoIncrement, - finalVectorDimension); + return new SnowflakeColumnMetadata(colNode, jdbcTreatDecimalAsInt, session); } static ColumnTypeInfo getSnowflakeType( @@ -443,11 +387,11 @@ static List createFieldsMetadata( return fields; } - private static boolean isVectorType(String internalColumnTypeName) { + static boolean isVectorType(String internalColumnTypeName) { return internalColumnTypeName.equalsIgnoreCase("vector"); } - private static List getFieldMetadata( + static List getFieldMetadata( boolean jdbcTreatDecimalAsInt, String internalColumnTypeName, JsonNode node) throws SnowflakeSQLLoggedException { if (!node.path("fields").isEmpty()) { diff --git a/src/test/java/net/snowflake/client/jdbc/SnowflakeSerializableTest.java b/src/test/java/net/snowflake/client/jdbc/SnowflakeSerializableTest.java new file mode 100644 index 000000000..4cd2fa7e8 --- /dev/null +++ b/src/test/java/net/snowflake/client/jdbc/SnowflakeSerializableTest.java @@ -0,0 +1,381 @@ +package net.snowflake.client.jdbc; + +import static net.snowflake.client.jdbc.SnowflakeChunkDownloader.NoOpChunkDownloader; +import static net.snowflake.client.jdbc.SnowflakeResultSetSerializableV1.ChunkFileMetadata; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import java.util.HashMap; +import net.snowflake.client.core.ObjectMapperFactory; +import net.snowflake.client.core.QueryResultFormat; +import net.snowflake.client.core.SFBaseSession; +import net.snowflake.client.core.SFBaseStatement; +import net.snowflake.client.core.SFStatementType; +import org.junit.Test; + +public class SnowflakeSerializableTest { + + private static final ObjectMapper OBJECT_MAPPER = ObjectMapperFactory.getObjectMapper(); + private static final String STANDARD_SERIALIZABLE_V1_JSON_STRING = + "{\n" + + " \"data\": {\n" + + " \"parameters\": [\n" + + " {\n" + + " \"name\": \"CLIENT_PREFETCH_THREADS\",\n" + + " \"value\": 4\n" + + " },\n" + + " {\n" + + " \"name\": \"TIMESTAMP_OUTPUT_FORMAT\",\n" + + " \"value\": \"YYYY-MM-DD HH24:MI:SS.FF3 TZHTZM\"\n" + + " },\n" + + " {\n" + + " \"name\": \"CLIENT_RESULT_CHUNK_SIZE\",\n" + + " \"value\": 128\n" + + " }\n" + + " ],\n" + + " \"rowtype\": [\n" + + " {\n" + + " \"name\": \"1\",\n" + + " \"database\": \"some-db\",\n" + + " \"schema\": \"some-schema\",\n" + + " \"table\": \"some-table\",\n" + + " \"byteLength\": null,\n" + + " \"type\": \"fixed\",\n" + + " \"length\": 256,\n" + + " \"scale\": 0,\n" + + " \"precision\": 1,\n" + + " \"nullable\": false,\n" + + " \"collation\": null\n" + + " }\n" + + " ],\n" + + " \"rowset\": [\n" + + " [\n" + + " \"1\"\n" + + " ]\n" + + " ],\n" + + " \"qrmk\": \"ADCDEFGHIJdwadawYhiF81aC0wT0IU+NN8QtobPWCk=\",\n" + + " \"chunkHeaders\": {\n" + + " \"x-amz-server-side-encryption-customer-key-md5\": \"A2dDf2ff7HI8OCdsR3pK82g==\"\n" + + " },\n" + + " \"chunks\": [\n" + + " {\n" + + " \"url\": \"https://sfc-ds2-customer-stage.s3.us-west-2.amazonaws.com\",\n" + + " \"rowCount\": 756,\n" + + " \"uncompressedSize\": 312560,\n" + + " \"compressedSize\": 26828\n" + + " }\n" + + " ],\n" + + " \"total\": 1,\n" + + " \"returned\": 1,\n" + + " \"queryId\": \"01b341c1-0000-772f-0000-0004189328ca\",\n" + + " \"databaseProvider\": \"some-db-provider\",\n" + + " \"finalDatabaseName\": \"some-db\",\n" + + " \"finalSchemaName\": \"some-schema\",\n" + + " \"finalWarehouseName\": \"some-warehouse\",\n" + + " \"finalRoleName\": \"ENG_OPS_RL\",\n" + + " \"numberOfBinds\": 0,\n" + + " \"arrayBindSupported\": false,\n" + + " \"statementTypeId\": 4096,\n" + + " \"version\": 1,\n" + + " \"sendResultTime\": 1711499620154,\n" + + " \"queryResultFormat\": \"json\",\n" + + " \"queryContext\": {\n" + + " \"entries\": [\n" + + " {\n" + + " \"id\": 0,\n" + + " \"timestamp\": 55456940208204,\n" + + " \"priority\": 0\n" + + " }\n" + + " ]\n" + + " }\n" + + " },\n" + + " \"code\": null,\n" + + " \"message\": null,\n" + + " \"success\": true\n" + + "}"; + + public static final String RICH_RESULTS_SERIALIZABLE_V1_JSON_STRING = + "{\n" + + " \"data\": {\n" + + " \"parameters\": [\n" + + " {\n" + + " \"name\": \"CLIENT_PREFETCH_THREADS\",\n" + + " \"value\": 4\n" + + " },\n" + + " {\n" + + " \"name\": \"TIMESTAMP_OUTPUT_FORMAT\",\n" + + " \"value\": \"YYYY-MM-DD HH24:MI:SS.FF3 TZHTZM\"\n" + + " },\n" + + " {\n" + + " \"name\": \"CLIENT_RESULT_CHUNK_SIZE\",\n" + + " \"value\": 128\n" + + " }\n" + + " ],\n" + + " \"rowtype\": [\n" + + " {\n" + + " \"name\": \"1\",\n" + + " \"database\": \"some-db\",\n" + + " \"schema\": \"some-schema\",\n" + + " \"table\": \"some-table\",\n" + + " \"byteLength\": null,\n" + + " \"type\": \"fixed\",\n" + + " \"length\": 256,\n" + + " \"scale\": 0,\n" + + " \"precision\": 1,\n" + + " \"nullable\": false,\n" + + " \"collation\": null\n" + + " }\n" + + " ],\n" + + " \"rowset\": [\n" + + " [\n" + + " \"1\"\n" + + " ]\n" + + " ],\n" + + " \"qrmk\": \"ADCDEFGHIJdwadawYhiF81aC0wT0IU+NN8QtobPWCk=\",\n" + + " \"chunkHeaders\": {\n" + + " \"x-amz-server-side-encryption-customer-key-md5\": \"A2dDf2ff7HI8OCdsR3pK82g==\"\n" + + " },\n" + + " \"chunks\": [\n" + + " {\n" + + " \"url\": \"https://sfc-ds2-customer-stage.s3.us-west-2.amazonaws.com\",\n" + + " \"rowCount\": 756,\n" + + " \"uncompressedSize\": 312560,\n" + + " \"compressedSize\": 26828\n" + + " }\n" + + " ],\n" + + " \"total\": 1,\n" + + " \"returned\": 1,\n" + + " \"queryId\": \"01b341c1-0000-772f-0000-0004189328ca\",\n" + + " \"databaseProvider\": \"some-db-provider\",\n" + + " \"finalDatabaseName\": \"some-db\",\n" + + " \"finalSchemaName\": \"some-schema\",\n" + + " \"finalWarehouseName\": \"some-warehouse\",\n" + + " \"finalRoleName\": \"ENG_OPS_RL\",\n" + + " \"numberOfBinds\": 0,\n" + + " \"arrayBindSupported\": false,\n" + + " \"statementTypeId\": 4096,\n" + + " \"version\": 1,\n" + + " \"sendResultTime\": 1711499620154,\n" + + " \"queryResultFormat\": \"json\",\n" + + " \"queryContext\": {\n" + + " \"entries\": [\n" + + " {\n" + + " \"id\": 0,\n" + + " \"timestamp\": 55456940208204,\n" + + " \"priority\": 0\n" + + " }\n" + + " ]\n" + + " }\n" + + " },\n" + + " \"richResult\": {\n" + + " \"rowtype\": [\n" + + " {\n" + + " \"name\": \"LOWER_BOUND\",\n" + + " \"database\": \"TEMP\",\n" + + " \"schema\": \"PUBLIC\",\n" + + " \"table\": \"T_TEST\",\n" + + " \"precision\": null,\n" + + " \"byteLength\": 16777216,\n" + + " \"type\": \"fixed\",\n" + + " \"scale\": null,\n" + + " \"nullable\": true,\n" + + " \"collation\": null,\n" + + " \"length\": 16777216,\n" + + " \"columnIndexing\": 1\n" + + " },\n" + + " {\n" + + " \"name\": \"UPPER_BOUND\",\n" + + " \"database\": \"TEMP\",\n" + + " \"schema\": \"PUBLIC\",\n" + + " \"table\": \"T_TEST\",\n" + + " \"precision\": null,\n" + + " \"byteLength\": 16777216,\n" + + " \"type\": \"fixed\",\n" + + " \"scale\": null,\n" + + " \"nullable\": true,\n" + + " \"collation\": null,\n" + + " \"length\": 16777216,\n" + + " \"columnIndexing\": 1\n" + + " }\n" + + " ],\n" + + " \"rowset\": [\n" + + " [\n" + + " \"value1_lower\",\n" + + " \"value1_upper\"\n" + + " ],\n" + + " [\n" + + " \"value2_lower\",\n" + + " \"value2_upper\"\n" + + " ]\n" + + "\n" + + " ],\n" + + " \"qrmk\": \"ZXYADCDEFGHIJdwadawYhiF81aC0wT0IU+NN8QtobPWCk=\",\n" + + " \"chunkHeaders\": {\n" + + " \"x-amz-server-side-encryption-customer-key-md5\": \"f342lkkftyf7HI8OCdsR3pK82g==\"\n" + + " },\n" + + " \"chunks\": [\n" + + " {\n" + + " \"url\": \"https://sfc-ds2-customer-stage.s3.us-west-2.amazonaws.com/rich-res\",\n" + + " \"rowCount\": 756,\n" + + " \"uncompressedSize\": 312560,\n" + + " \"compressedSize\": 26828\n" + + " }\n" + + " ],\n" + + " \"total\": 1,\n" + + " \"queryId\": \"01b341c1-0000-772f-0000-0004189328ca\",\n" + + " \"queryResultFormat\": \"json\"\n" + + " },\n" + + " \"code\": null,\n" + + " \"message\": null,\n" + + " \"success\": true\n" + + "}"; + + private static final SFBaseSession MOCK_SESSION = + new MockConnectionTest.MockSnowflakeConnectionImpl().getSFSession(); + private static final SFBaseStatement MOCK_STATEMENT = + new MockConnectionTest.MockSnowflakeConnectionImpl().getSFStatement(); + + @Test + public void shouldProperlyCreateSerializableV1() + throws JsonProcessingException, SnowflakeSQLException { + JsonNode rootNode = OBJECT_MAPPER.readTree(STANDARD_SERIALIZABLE_V1_JSON_STRING); + SnowflakeResultSetSerializableV1 s = + SnowflakeResultSetSerializableV1.create( + rootNode, MOCK_SESSION, MOCK_STATEMENT, new DefaultResultStreamProvider()); + assertRegularResultSetSerializable(s, SnowflakeChunkDownloader.class); + } + + @Test + public void shouldCreateSerializableWithNoOpChunksDownloader() + throws JsonProcessingException, SnowflakeSQLException { + JsonNode rootNode = OBJECT_MAPPER.readTree(STANDARD_SERIALIZABLE_V1_JSON_STRING); + SnowflakeResultSetSerializableV1 s = + SnowflakeResultSetSerializableV1.createWithChunksPrefetchDisabled( + rootNode, MOCK_SESSION, MOCK_STATEMENT); + assertRegularResultSetSerializable(s, NoOpChunkDownloader.class); + } + + @Test + public void shouldProperlyCreateRichSerializableV1() + throws JsonProcessingException, SnowflakeSQLException { + JsonNode rootNode = OBJECT_MAPPER.readTree(RICH_RESULTS_SERIALIZABLE_V1_JSON_STRING); + SnowflakeRichResultSetSerializableV1 s = + SnowflakeRichResultSetSerializableV1.createWithChunksPrefetchDisabled( + rootNode, MOCK_SESSION, MOCK_STATEMENT); + assertRegularResultSetSerializable(s, NoOpChunkDownloader.class); + assertRichResultSetSerializable(s); + } + + private void assertRegularResultSetSerializable( + SnowflakeResultSetSerializableV1 s, Class expectedChunkDownloaderType) { + assertNotNull(s); + assertEquals("01b341c1-0000-772f-0000-0004189328ca", s.getQueryId()); + assertEquals("some-db", s.getFinalDatabaseName()); + assertEquals("some-schema", s.getFinalSchemaName()); + assertEquals("some-warehouse", s.getFinalWarehouseName()); + assertEquals("ENG_OPS_RL", s.getFinalRoleName()); + assertEquals(0, s.getNumberOfBinds()); + assertEquals(QueryResultFormat.JSON, s.getQueryResultFormat()); + assertEquals(SFStatementType.SELECT, s.getStatementType()); + assertEquals( + new HashMap() { + { + put("CLIENT_PREFETCH_THREADS", 4); + put("TIMESTAMP_OUTPUT_FORMAT", "YYYY-MM-DD HH24:MI:SS.FF3 TZHTZM"); + put("CLIENT_RESULT_CHUNK_SIZE", 128); + } + }, + s.getParameters()); + assertEquals("ADCDEFGHIJdwadawYhiF81aC0wT0IU+NN8QtobPWCk=", s.getQrmk()); + assertFalse(s.isArrayBindSupported()); + assertEquals(1, s.getResultVersion()); + + // column metadata + assertEquals(1, s.getColumnCount()); + assertEquals(1, s.getResultColumnMetadata().size()); + SnowflakeColumnMetadata column = s.getResultColumnMetadata().get(0); + assertEquals("1", column.getName()); + assertEquals("NUMBER", column.getTypeName()); + assertEquals(-5, column.getType()); + assertFalse(column.isNullable()); + assertEquals(256, column.getLength()); + assertEquals(0, column.getScale()); + assertEquals(1, column.getPrecision()); + assertEquals(SnowflakeType.FIXED, column.getBase()); + assertEquals("some-db", column.getColumnSrcDatabase()); + assertEquals("some-schema", column.getColumnSrcSchema()); + assertEquals("some-table", column.getColumnSrcTable()); + + // chunks metadata + assertEquals("[[\"1\"]]", s.getFirstChunkStringData()); + assertEquals(1, s.getChunkHeadersMap().size()); + assertEquals( + "A2dDf2ff7HI8OCdsR3pK82g==", + s.getChunkHeadersMap().get("x-amz-server-side-encryption-customer-key-md5")); + assertEquals(1, s.getChunkFileCount()); + assertEquals(1, s.getChunkFileMetadatas().size()); + ChunkFileMetadata chunkMeta = s.getChunkFileMetadatas().get(0); + assertEquals(756, chunkMeta.getRowCount()); + assertEquals(26828, chunkMeta.getCompressedByteSize()); + assertEquals(312560, chunkMeta.getUncompressedByteSize()); + assertEquals( + "https://sfc-ds2-customer-stage.s3.us-west-2.amazonaws.com", chunkMeta.getFileURL()); + assertNotNull(s.chunkDownloader); + assertTrue(expectedChunkDownloaderType.isInstance(s.chunkDownloader)); + } + + private void assertRichResultSetSerializable(SnowflakeRichResultSetSerializableV1 s) { + // column metadata + assertEquals(2, s.getRichResultsColumnCount()); + assertEquals(2, s.getRichResultsColumnMetadata().size()); + SnowflakeRichResultSetSerializableV1.SnowflakeRichResultsColumnMetadata lowerBound = + s.getRichResultsColumnMetadata().get(0); + assertEquals("LOWER_BOUND", lowerBound.getName()); + assertEquals("NUMBER", lowerBound.getTypeName()); + assertEquals(-5, lowerBound.getType()); + assertTrue(lowerBound.isNullable()); + assertEquals(16777216, lowerBound.getLength()); + assertEquals(SnowflakeType.FIXED, lowerBound.getBase()); + assertEquals("TEMP", lowerBound.getColumnSrcDatabase()); + assertEquals("PUBLIC", lowerBound.getColumnSrcSchema()); + assertEquals("T_TEST", lowerBound.getColumnSrcTable()); + assertEquals(1, lowerBound.getColumnIndex()); + + SnowflakeRichResultSetSerializableV1.SnowflakeRichResultsColumnMetadata upperBound = + s.getRichResultsColumnMetadata().get(1); + assertEquals("UPPER_BOUND", upperBound.getName()); + assertEquals("NUMBER", upperBound.getTypeName()); + assertEquals(-5, upperBound.getType()); + assertTrue(upperBound.isNullable()); + assertEquals(16777216, upperBound.getLength()); + assertEquals(SnowflakeType.FIXED, upperBound.getBase()); + assertEquals("TEMP", upperBound.getColumnSrcDatabase()); + assertEquals("PUBLIC", upperBound.getColumnSrcSchema()); + assertEquals("T_TEST", upperBound.getColumnSrcTable()); + assertEquals(1, upperBound.getColumnIndex()); + + // chunks metadata + assertEquals( + "[[\"value1_lower\",\"value1_upper\"],[\"value2_lower\",\"value2_upper\"]]", + s.getRichResultsFirstChunkStringData()); + assertEquals("ZXYADCDEFGHIJdwadawYhiF81aC0wT0IU+NN8QtobPWCk=", s.getRichResultsQrmk()); + assertEquals(1, s.getRichResultsChunkHeadersMap().size()); + assertEquals( + "f342lkkftyf7HI8OCdsR3pK82g==", + s.getRichResultsChunkHeadersMap().get("x-amz-server-side-encryption-customer-key-md5")); + assertEquals(1, s.getRichResultsChunkFileCount()); + assertEquals(1, s.getRichResultsChunkFilesMetadata().size()); + ChunkFileMetadata chunkMeta = s.getRichResultsChunkFilesMetadata().get(0); + assertEquals(756, chunkMeta.getRowCount()); + assertEquals(26828, chunkMeta.getCompressedByteSize()); + assertEquals(312560, chunkMeta.getUncompressedByteSize()); + assertEquals( + "https://sfc-ds2-customer-stage.s3.us-west-2.amazonaws.com/rich-res", + chunkMeta.getFileURL()); + } +}