Snowflake Connector - Reduce computing resources used for metadata queries #43452
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -293,4 +293,8 @@ object SnowflakeDatabaseUtils { | |
| AirbyteProtocolType.UNKNOWN -> "VARIANT" | ||
| } | ||
| } | ||
|
|
||
| fun fromIsNullableSnowflakeString(isNullable: String?): Boolean { | ||
| return "true".equals(isNullable, ignoreCase = true) | ||
|
There was a problem hiding this comment. remove and use String.toBoolean() There was a problem hiding this comment. Removed the function and used String.toBoolean |
||
| } | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -26,17 +26,18 @@ import io.airbyte.integrations.base.destination.typing_deduping.Struct | |
| import io.airbyte.integrations.base.destination.typing_deduping.Union | ||
| import io.airbyte.integrations.base.destination.typing_deduping.UnsupportedOneOf | ||
| import io.airbyte.integrations.destination.snowflake.SnowflakeDatabaseUtils | ||
| import io.airbyte.integrations.destination.snowflake.SnowflakeDatabaseUtils.fromIsNullableSnowflakeString | ||
| import io.airbyte.integrations.destination.snowflake.migrations.SnowflakeState | ||
| import io.airbyte.integrations.destination.snowflake.typing_deduping.SnowflakeSqlGenerator.Companion.QUOTE | ||
| import java.sql.Connection | ||
| import java.sql.ResultSet | ||
| import java.sql.SQLException | ||
| import java.time.Instant | ||
| import java.util.* | ||
| import java.util.stream.Collectors | ||
| import net.snowflake.client.jdbc.SnowflakeSQLException | ||
| import org.apache.commons.text.StringSubstitutor | ||
| import org.json.JSONObject | ||
| import org.jooq.SQLDialect | ||
| import org.slf4j.Logger | ||
| import org.slf4j.LoggerFactory | ||
|
|
@@ -70,75 +71,83 @@ class SnowflakeDestinationHandler( | |
| private fun getFinalTableRowCount( | ||
| streamIds: List<StreamId> | ||
| ): LinkedHashMap<String, LinkedHashMap<String, Int>> { | ||
|
|
||
| val tableRowCountsFromShowQuery = LinkedHashMap<String, LinkedHashMap<String, Int>>() | ||
| var showColumnsResult: List<JsonNode> = listOf() | ||
|
There was a problem hiding this comment. remove this and set the val inside your try block |
||
|
|
||
| try { | ||
| for (stream in streamIds) { | ||
| val showColumnsQuery = | ||
| String.format( | ||
| """ | ||
|
There was a problem hiding this comment. use kotlin templates instead of |
||
| SHOW TABLES LIKE '%s' IN "%s"."%s"; | ||
| """.trimIndent(), | ||
| stream.finalName, | ||
| databaseName, | ||
| stream.finalNamespace, | ||
| ) | ||
| showColumnsResult = database.queryJsons( | ||
| showColumnsQuery, | ||
| ) | ||
| for (result in showColumnsResult) { | ||
| val tableSchema = result["schema_name"].asText() | ||
| val tableName = result["name"].asText() | ||
| val rowCount = result["rows"].asText() | ||
|
|
||
| tableRowCountsFromShowQuery | ||
|
There was a problem hiding this comment. The indentation is super confusing here (probably enforced by our format command). Any way to change that, or is our formatter going to bark at you? There was a problem hiding this comment.
|
||
| .computeIfAbsent(tableSchema) { _: String? -> LinkedHashMap() }[tableName] = | ||
| rowCount.toInt() | ||
| } | ||
| } | ||
| } catch (e: SQLException) { | ||
| showColumnsResult.stream().close() | ||
|
There was a problem hiding this comment. not sure why we need to close the stream here |
||
| //Not re-throwing the exception since the SQLException occurs when the table does not exist | ||
| //throw e | ||
|
There was a problem hiding this comment. remove the commented |
||
| } | ||
| return tableRowCountsFromShowQuery | ||
| } | ||
|
|
||
|
|
||
| @Throws(Exception::class) | ||
| private fun getInitialRawTableState( | ||
| id: StreamId, | ||
| suffix: String, | ||
| ): InitialRawTableStatus { | ||
|
|
||
| val rawTableName = id.rawName + suffix | ||
| var tableExists = false | ||
| var showTablesResult: List<JsonNode> = listOf() | ||
|
|
||
| try { | ||
| val showTablesQuery = | ||
| String.format( | ||
| """ | ||
| SHOW TABLES LIKE '%s' IN "%s"."%s"; | ||
|
There was a problem hiding this comment. use kotlin templates instead of There was a problem hiding this comment. does that work when the QUOTED_IDENTIFIERS_IGNORE_CASE is set to true? We have a test class that sets that for the testDatabase |
||
| """.trimIndent(), | ||
| rawTableName, | ||
| databaseName, | ||
| id.rawNamespace, | ||
| ) | ||
| showTablesResult = database.queryJsons( | ||
| showTablesQuery, | ||
| ) | ||
| if(showTablesResult.size > 0) { | ||
| tableExists = true | ||
| } | ||
| } catch (e: SQLException) { | ||
| showTablesResult.stream().close() | ||
| //Not re-throwing the exception since the SQLException occurs when the table does not exist | ||
| //throw e | ||
| } | ||
|
|
||
| if (!tableExists) { | ||
| return InitialRawTableStatus( | ||
| rawTableExists = false, | ||
| hasUnprocessedRecords = false, | ||
| maxProcessedTimestamp = Optional.empty() | ||
| ) | ||
| } | ||
|
|
||
| // Snowflake timestamps have nanosecond precision, so decrement by 1ns | ||
| // And use two explicit queries because COALESCE doesn't short-circuit. | ||
| // This first query tries to find the oldest raw record with loaded_at = NULL | ||
|
|
@@ -560,6 +569,7 @@ class SnowflakeDestinationHandler( | |
| } | ||
|
|
||
| companion object { | ||
|
|
||
| private val LOGGER: Logger = | ||
| LoggerFactory.getLogger(SnowflakeDestinationHandler::class.java) | ||
| const val EXCEPTION_COMMON_PREFIX: String = | ||
|
|
@@ -573,39 +583,63 @@ class SnowflakeDestinationHandler( | |
| databaseName: String, | ||
| streamIds: List<StreamId> | ||
| ): LinkedHashMap<String, LinkedHashMap<String, TableDefinition>> { | ||
|
|
||
| val existingTablesFromShowQuery = | ||
| LinkedHashMap<String, LinkedHashMap<String, TableDefinition>>() | ||
| var showColumnsResult: List<JsonNode> = listOf() | ||
|
|
||
| try { | ||
| for (stream in streamIds) { | ||
| val showColumnsQuery = | ||
| String.format( | ||
| """ | ||
| SHOW COLUMNS IN TABLE "%s"."%s"."%s"; | ||
| """.trimIndent(), | ||
| databaseName, | ||
| stream.finalNamespace, | ||
| stream.finalName, | ||
| ) | ||
| showColumnsResult = database.queryJsons( | ||
| showColumnsQuery, | ||
| ) | ||
|
|
||
| for (result in showColumnsResult) { | ||
| val tableSchema = result["schema_name"].asText() | ||
| val tableName = result["table_name"].asText() | ||
| val columnName = result["column_name"].asText() | ||
| var dataType = JSONObject(result["data_type"].asText()).getString("type") | ||
|
|
||
| //TODO: Need to check if there are other datatype differences | ||
|
There was a problem hiding this comment. that should probably be checked before merging this PR. Do we have any automated test that checks this? There was a problem hiding this comment. per the doc, it seems those are the only differences. I'd rather have a comment that points to https://docs.snowflake.com/en/sql-reference/sql/show-columns#output rather than a |
||
| // between the original approach and the new approach with SHOW queries | ||
| if(dataType.equals("FIXED")) { | ||
| dataType = "NUMBER" | ||
| } else if(dataType.equals("REAL")) { | ||
| dataType = "FLOAT" | ||
| } | ||
|
|
||
| val isNullable = result["null?"].asText() | ||
| val tableDefinition = | ||
|
There was a problem hiding this comment. let's create the table definition outside of the loop on |
||
| existingTablesFromShowQuery | ||
| .computeIfAbsent(tableSchema) { _: String? -> LinkedHashMap() } | ||
| .computeIfAbsent(tableName) { _: String? -> | ||
| TableDefinition(LinkedHashMap()) | ||
| } | ||
| tableDefinition.columns[columnName] = | ||
| ColumnDefinition( | ||
| columnName, | ||
| dataType, | ||
| 0, | ||
| fromIsNullableSnowflakeString(isNullable), | ||
| ) | ||
| } | ||
| } | ||
| } catch (e: SQLException) { | ||
| showColumnsResult.stream().close() | ||
| //Not re-throwing the exception since the SQLException occurs when the table does not exist | ||
| //throw e | ||
| } | ||
| return existingTablesFromShowQuery | ||
| } | ||
| } | ||
| } | ||
|
|
||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
connection.use. Also, keep connection a val instead of a var
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sure, I have changed the connection to a val. The code won't be changed to connection.use as we discussed since the connection needs to be open when the result set is returned