def test_stream_name_to_dict(self):
"""Test identifying catalog, schema and table names from fully qualified stream and table names"""
# Singer stream name format (Default '-' separator)
self.assertEqual(stream_utils.stream_name_to_dict('my_table'),
{"catalog_name": None, "schema_name": None, "table_name": "my_table"})
# Singer stream name format (Default '-' separator)
self.assertEqual(stream_utils.stream_name_to_dict('my_schema-my_table'),
{"catalog_name": None, "schema_name": "my_schema", "table_name": "my_table"})
# Singer stream name format (Default '-' separator)
self.assertEqual(stream_utils.stream_name_to_dict('my_catalog-my_schema-my_table'),
{"catalog_name": "my_catalog", "schema_name": "my_schema", "table_name": "my_table"})
# Snowflake table format (Custom '.' separator)
self.assertEqual(stream_utils.stream_name_to_dict('my_table', separator='.'),
{"catalog_name": None, "schema_name": None, "table_name": "my_table"})
# Snowflake table format (Custom '.' separator)
self.assertEqual(stream_utils.stream_name_to_dict('my_schema.my_table', separator='.'),
{"catalog_name": None, "schema_name": "my_schema", "table_name": "my_table"})
# Snowflake table format (Custom '.' separator)
self.assertEqual(stream_utils.stream_name_to_dict('my_catalog.my_schema.my_table', separator='.'),
{"catalog_name": "my_catalog", "schema_name": "my_schema", "table_name": "my_table"})
def table_name(self, stream_name, is_temporary, without_schema=False):
"""Generate target table name"""
if not stream_name:
return None
stream_dict = stream_utils.stream_name_to_dict(stream_name)
table_name = stream_dict['table_name']
sf_table_name = table_name.replace('.', '_').replace('-', '_').lower()
if is_temporary:
sf_table_name = '{}_temp'.format(sf_table_name)
if without_schema:
return f'"{sf_table_name.upper()}"'
return f'{self.schema_name}."{sf_table_name.upper()}"'
def __init__(self, connection_config, stream_schema_message=None, table_cache=None, file_format_type=None):
"""
connection_config: Snowflake connection details
stream_schema_message: An instance of the DbSync class is typically used to load
data only from a certain singer tap stream.
The stream_schema_message holds the destination schema
name and the JSON schema that will be used to
validate every RECORDS messages that comes from the stream.
Schema validation happening before creating CSV and before
uploading data into Snowflake.
If stream_schema_message is not defined that we can use
the DbSync instance as a generic purpose connection to
Snowflake and can run individual queries. For example
collecting catalog informations from Snowflake for caching
purposes.
"""
self.connection_config = connection_config
self.stream_schema_message = stream_schema_message
self.table_cache = table_cache
# logger to be used across the class's methods
self.logger = get_logger('target_snowflake')
# Validate connection configuration
config_errors = validate_config(connection_config)
# Exit if config has errors
if len(config_errors) > 0:
self.logger.error('Invalid configuration:\n * %s', '\n * '.join(config_errors))
sys.exit(1)
if self.connection_config.get('stage', None):
stage = stream_utils.stream_name_to_dict(self.connection_config['stage'], separator='.')
if not stage['schema_name']:
self.logger.error(
"The named external stage object in config has to use the <schema>.<stage_name> format.")
sys.exit(1)
self.schema_name = None
self.grantees = None
self.file_format = FileFormat(self.connection_config['file_format'], self.query, file_format_type)
if not self.connection_config.get('stage') and self.file_format.file_format_type == FileFormatTypes.PARQUET:
self.logger.error("Table stages with Parquet file format is not suppported. "
"Use named stages with Parquet file format or table stages with CSV files format")
sys.exit(1)
# Init stream schema pylint: disable=line-too-long
if self.stream_schema_message is not None:
# Define target schema name.
# --------------------------
# Target schema name can be defined in multiple ways:
#
# 1: 'default_target_schema' key : Target schema is the same for every incoming stream if
# not specified explicitly for a given stream in
# the `schema_mapping` object
# 2: 'schema_mapping' key : Target schema defined explicitly for a given stream.
# Example config.json:
# "schema_mapping": {
# "my_tap_stream_id": {
# "target_schema": "my_snowflake_schema",
# "target_schema_select_permissions": [ "role_with_select_privs" ]
# }
# }
config_default_target_schema = self.connection_config.get('default_target_schema', '').strip()
config_schema_mapping = self.connection_config.get('schema_mapping', {})
stream_name = stream_schema_message['stream']
stream_schema_name = stream_utils.stream_name_to_dict(stream_name)['schema_name']
if config_schema_mapping and stream_schema_name in config_schema_mapping:
self.schema_name = config_schema_mapping[stream_schema_name].get('target_schema')
elif config_default_target_schema:
self.schema_name = config_default_target_schema
if not self.schema_name:
raise Exception(
"Target schema name not defined in config. "
"Neither 'default_target_schema' (string) nor 'schema_mapping' (object) defines "
"target schema for {} stream.".format(stream_name))
# Define grantees
# ---------------
# Grantees can be defined in multiple ways:
#
# 1: 'default_target_schema_select_permissions' key : USAGE and SELECT privileges will be granted on every table to a given role
# for every incoming stream if not specified explicitly
# in the `schema_mapping` object
# 2: 'target_schema_select_permissions' key : Roles to grant USAGE and SELECT privileges defined explicitly
# for a given stream.
# Example config.json:
# "schema_mapping": {
# "my_tap_stream_id": {
# "target_schema": "my_snowflake_schema",
# "target_schema_select_permissions": [ "role_with_select_privs" ]
# }
# }
self.grantees = self.connection_config.get('default_target_schema_select_permissions')
if config_schema_mapping and stream_schema_name in config_schema_mapping:
self.grantees = config_schema_mapping[stream_schema_name].get('target_schema_select_permissions',
self.grantees)
self.data_flattening_max_level = self.connection_config.get('data_flattening_max_level', 0)
self.flatten_schema = flattening.flatten_schema(stream_schema_message['schema'],
max_level=self.data_flattening_max_level)
# Use external stage
if connection_config.get('s3_bucket', None):
self.upload_client = S3UploadClient(connection_config)
# Use table stage
else:
# Enforce no parallelism with table stages.
# The PUT command in the snowflake-python-connector is using boto3.create_client() function which is not
# thread safe. More info at https://github.com/boto/boto3/issues/801
if connection_config.get('parallelism') != 1:
self.logger.warning('Enforcing to use single thread parallelism with table stages. '
'The PUT command in the snowflake-python-connector is using boto3 create_client() '
'function which is not thread safe. '
'If you need parallel file upload please use external stage by adding s3_bucket '
'key to the configuration')
connection_config['parallelism'] = 1
self.upload_client = SnowflakeUploadClient(connection_config, self)
def flush_records(stream: str,
records: List[Dict],
db_sync: DbSync,
temp_dir: str = None,
no_compression: bool = False,
archive_load_files: Dict = None) -> None:
"""
Takes a list of record messages and loads it into the snowflake target table
Args:
stream: Name of the stream
records: List of dictionary, that represents multiple csv lines. Dict key is the column name, value is the
column value
row_count:
db_sync: A DbSync object
temp_dir: Directory where intermediate temporary files will be created. (Default: OS specific temp directory)
no_compression: Disable to use compressed files. (Default: False)
archive_load_files: Data needed for archive load files. (Default: None)
Returns:
None
"""
# Generate file on disk in the required format
filepath = db_sync.file_format.formatter.records_to_file(records,
db_sync.flatten_schema,
compression=not no_compression,
dest_dir=temp_dir,
data_flattening_max_level=
db_sync.data_flattening_max_level)
# Get file stats
row_count = len(records)
size_bytes = os.path.getsize(filepath)
# Upload to s3 and load into Snowflake
s3_key = db_sync.put_to_stage(filepath, stream, row_count, temp_dir=temp_dir)
db_sync.load_file(s3_key, row_count, size_bytes)
# Delete file from local disk
os.remove(filepath)
if archive_load_files:
stream_name_parts = stream_utils.stream_name_to_dict(stream)
if 'schema_name' not in stream_name_parts or 'table_name' not in stream_name_parts:
raise Exception("Failed to extract schema and table names from stream '{}'".format(stream))
archive_schema = stream_name_parts['schema_name']
archive_table = stream_name_parts['table_name']
archive_tap = archive_load_files['tap']
archive_metadata = {
'tap': archive_tap,
'schema': archive_schema,
'table': archive_table,
'archived-by': 'pipelinewise_target_snowflake'
}
if 'column' in archive_load_files:
archive_metadata.update({
'incremental-key': archive_load_files['column'],
'incremental-key-min': str(archive_load_files['min']),
'incremental-key-max': str(archive_load_files['max'])
})
# Use same file name as in import
archive_file = os.path.basename(s3_key)
archive_key = "{}/{}/{}".format(archive_tap, archive_table, archive_file)
db_sync.copy_to_archive(s3_key, archive_key, archive_metadata)
# Delete file from S3
db_sync.delete_from_stage(stream, s3_key)