"""Helpers for the s32sf workflow's tasks.""" import re import boto3 from garcon_contrib.aws.utils import garcon_s3 from snowflake_etl.conf import config as cfg class SourcesConf: """Lightweight class to store attributes of the Redshift source tables.""" def __init__(self): """Create a SourceConf entity.""" schemas_conf = {} for schema, tables in schemas_conf.items(): setattr(self, schema, tables) def is_table_transient(self, table, schema): """Check if we need to create stg and dest tables as transient. If table wasn't found in the sources.yml within a specified schema defaults to not transient. Args: table (str): Table name. schema (str): Schema name. Returns: str: Table param to format a SQL statement. '' or 'TRANSIENT'. Raises: ValueError: If table wasn't found in the sources.yml within a specified schema. """ # in case if we're taking schema name from Snowflake config schema = schema.lower() tables = self.__getattribute__(schema) table_attrs = tables.get(table) if table_attrs: if table_attrs.get('transient'): return 'TRANSIENT' else: return '' else: # no table found in Snowflake config, default to not transient return '' def read_json_schema_on_s3(s3_path, logger=None): """Retrieve content of JSON schema file. Retrieve and output the content of JSON file which consists of the schema of Redshift table. Args: s3_path (str): S3 path to the JSON file. Returns: dict: JSON table schema. """ bucket, absolute_path = garcon_s3.extract_bucket_path(s3_path) # FIXME clean up - temporary measure - this whole function should be a # generic task in garcon contrib for S3 if logger: logger.info('Getting contents from S3 %s/%s', bucket, absolute_path) s3 = boto3.client('s3') s3_object = s3.get_object(Bucket=bucket, Key=absolute_path) return s3_object['Body'].read().decode('utf-8') # FIXME refactor this and underlying systems to just take in the partial SQL # for STAGE_FILE_FORMAT or be complete with wrapping all possible options def create_replace_temp_table_sql( db, schema, table, s3_path_data, file_type, field_delimiter, enclosed_by, compression): """Build create replace temp table SQL. Args: db (str): Name of the Snowflake database. schema (str): Name of the Snowflake schema. table (str): Name of the Snowflake table. s3_path_data (str): Full S3 path to data source. file_type (str): File type. field_delimiter (str): Delimiter of the S3 file. enclosed_by (str): Enclosing char of the S3 file. compression (str): Compression format of S3 file. Return: str: SQL for create replace Snowflake table. """ staging_table_name = '{db}.{schema}.staging_{table}'.format( db=db, schema=schema, table=table) full_table_name = '{db}.{schema}.{table}'.format( db=db, schema=schema, table=table) sql = cfg.SF_DEFAULT_QUERIES['create_replace_table_like'].format( full_table_name=full_table_name, staging_table_name=staging_table_name, s3_path_data=s3_path_data, aws_key_id=cfg.getconf('aws')['aws']['access_key'], aws_secret_key=cfg.getconf('aws')['aws']['access_secret'], file_type=file_type, field_delimiter=field_delimiter, enclosed_by=enclosed_by, compression=compression) return ' '.join(sql.split()) def validate_file_format(file_format): """Validate FILE_FORMAT Snowflake statement. Args: file_format (list): FILE_FORMAT statement from initial context. Raises: ValueError: If validation fails. """ for option in file_format: if re.match(r'^SKIP_HEADER=\d+$', option): continue elif (re.match(r'^DATE_FORMAT="\w{4}-\w{2}-\w{2}"$', option) or option == 'DATE_FORMAT="AUTO"'): continue elif (re.match( r'^TIMESTAMP_FORMAT="\w{4}-\w{2}-\w{2}\s\w{2}:\w{2}:' r'\w{2}\.\w{6}"$', option) or re.match( r'^TIMESTAMP_FORMAT="\w{4}-\w{2}-\w{2}T\w{2}:\w{2}:' r'\w{2}"$', option) or option == 'TIMESTAMP_FORMAT="AUTO"'): continue elif re.match(r'^TYPE="(CSV|AVRO|XML|JSON)"$', option): continue elif re.match( r'^COMPRESSION="(AUTO|GZIP|BZ2|DEFLATE|RAW_DEFLATE|NONE)"$', option): continue elif re.match(r'^ESCAPE="\\134"$', option): continue elif re.match(r'^ESCAPE_UNENCLOSED_FIELD="\\134"$', option): continue elif re.match(r'^RECORD_DELIMITER="(\\n|\\r|NONE)"$', option): continue elif re.match(r'^FIELD_DELIMITER="(,|;|\||\\t|\s|NONE)"$', option): continue elif re.match( r'^NULL_IF=(\("__NULL__"\)|\("\w{4}-\w{2}-\w{2}"\)|' r'\("NULL"\)|\(\))$', option): continue elif re.match( r'^FIELD_OPTIONALLY_ENCLOSED_BY="(\'|\"|NONE)"$', option): continue elif re.match(r"^FIELD_OPTIONALLY_ENCLOSED_BY='(\"|NONE)'$", option): continue elif re.match(r'^TRIM_SPACE=(TRUE|FALSE)$', option): continue elif re.match( r'^ERROR_ON_COLUMN_COUNT_MISMATCH=(TRUE|FALSE)$', option): continue elif re.match(r'^ENABLE_OCTAL=(TRUE|FALSE)$', option): continue elif re.match(r'^ALLOW_DUPLICATE=(TRUE|FALSE)$', option): continue elif re.match(r'^STRIP_OUTER_ARRAY=(TRUE|FALSE)$', option): continue elif re.match(r'^STRIP_NULL_VALUES=(TRUE|FALSE)$', option): continue elif re.match(r'^IGNORE_UTF8_ERRORS=(TRUE|FALSE)$', option): continue elif re.match(r'^PRESERVE_SPACE=(TRUE|FALSE)$', option): continue elif re.match(r'^STRIP_OUTER_ELEMENT=(TRUE|FALSE)$', option): continue elif re.match(r'^DISABLE_SNOWFLAKE_DATA=(TRUE|FALSE)$', option): continue elif re.match(r'^DISABLE_AUTO_CONVERT=(TRUE|FALSE)$', option): continue else: raise ValueError('Passed FILE_FORMAT not valid') def extract_db_and_schema(sfdb_params, db_key='db', schema_key='schema'): """Extract db name and db schema. Extracts db name and db schema from a dict with SFDB configuration parameters. Args: sfdb_params (dict): a dict of SFDB credentials and other parameters db_key (str), optional: a key for db name value schema_key (str), optional: a key for db schema value Returns: db, schema (tuple): If validation fails. """ db = sfdb_params.get(db_key) schema = sfdb_params.get(schema_key) return db, schema