AP-1011 archive_load_files feature

.circleci/config.yml

@@ -25,7 +25,7 @@ jobs:
  command: |
  . venv/bin/activate
  export LOGGING_CONF_FILE=$(pwd)/sample_logging.conf
- pytest tests/unit -vv --cov target_snowflake --cov-fail-under=55
+ pytest tests/unit -vv --cov target_snowflake --cov-fail-under=60
 
  - run:
  name: 'Integration Tests'

README.md

@@ -176,6 +176,7 @@ Full list of options in `config.json`:
 | temp_dir | String | | (Default: platform-dependent) Directory of temporary files with RECORD messages. |
 | no_compression | Boolean | | (Default: False) Generate uncompressed files when loading to Snowflake. Normally, by default GZIP compressed files are generated. |
 | query_tag | String | | (Default: None) Optional string to tag executed queries in Snowflake. Replaces tokens `{{database}}`, `{{schema}}` and `{{table}}` with the appropriate values. The tags are displayed in the output of the Snowflake `QUERY_HISTORY`, `QUERY_HISTORY_BY_*` functions. |
+| archive_load_files | Boolean | | (Default: False) When enabled, the files loaded to Snowflake will also be stored in `{{s3_bucket}}` under the key `/archive/{schema_name}/{table_name}/`. All archived files will have `tap`, `schema`, `table` and `archived-by` as S3 metadata keys. When incremental replication is used, the archived files will also have the following S3 metadata keys: `incremental-key`, `incremental-key-min` and `incremental-key-max`. 
 
 ### To run tests:
 

target_snowflake/__init__.py

@@ -74,7 +74,7 @@ def get_snowflake_statics(config):
  if not ('disable_table_cache' in config and config['disable_table_cache']):
  LOGGER.info('Getting catalog objects from table cache...')
 
- db = DbSync(config) # pylint: disable=invalid-name
+ db = DbSync(config)  # pylint: disable=invalid-name
  table_cache = db.get_table_columns(
  table_schemas=stream_utils.get_schema_names_from_config(config))
 
@@ -83,8 +83,9 @@ def get_snowflake_statics(config):
 
  return table_cache, file_format_type
 
+
 # pylint: disable=too-many-locals,too-many-branches,too-many-statements,invalid-name
-def persist_lines(config, lines, table_cache=None, file_format_type: FileFormatTypes=None) -> None:
+def persist_lines(config, lines, table_cache=None, file_format_type: FileFormatTypes = None) -> None:
  """Main loop to read and consume singer messages from stdin
 
  Params:
@@ -113,6 +114,8 @@ def persist_lines(config, lines, table_cache=None, file_format_type: FileFormatT
  batch_size_rows = config.get('batch_size_rows', DEFAULT_BATCH_SIZE_ROWS)
  batch_wait_limit_seconds = config.get('batch_wait_limit_seconds', None)
  flush_timestamp = datetime.utcnow()
+ archive_load_files = config.get('archive_load_files', False)
+ archive_load_files_data = {}
 
  # Loop over lines from stdin
  for line in lines:
@@ -170,6 +173,21 @@ def persist_lines(config, lines, table_cache=None, file_format_type: FileFormatT
  else:
  records_to_load[stream][primary_key_string] = o['record']
 
+ if archive_load_files and stream in archive_load_files_data:
+ # Keep track of min and max of the designated column
+ stream_archive_load_files_values = archive_load_files_data[stream]
+ if 'column' in stream_archive_load_files_values:
+ incremental_key_column_name = stream_archive_load_files_values['column']
+ incremental_key_value = o['record'][incremental_key_column_name]
+ min_value = stream_archive_load_files_values['min']
+ max_value = stream_archive_load_files_values['max']
+
+ if min_value is None or min_value > incremental_key_value:
+ stream_archive_load_files_values['min'] = incremental_key_value
+
+ if max_value is None or max_value < incremental_key_value:
+ stream_archive_load_files_values['max'] = incremental_key_value
+
  flush = False
  if row_count[stream] >= batch_size_rows:
  flush = True
@@ -196,6 +214,7 @@ def persist_lines(config, lines, table_cache=None, file_format_type: FileFormatT
  config,
  state,
  flushed_state,
+ archive_load_files_data,
  filter_streams=filter_streams)
 
  flush_timestamp = datetime.utcnow()
@@ -226,7 +245,8 @@ def persist_lines(config, lines, table_cache=None, file_format_type: FileFormatT
  stream_to_sync,
  config,
  state,
- flushed_state)
+ flushed_state,
+ archive_load_files_data)
 
  # emit latest encountered state
  emit_state(flushed_state)
@@ -257,6 +277,27 @@ def persist_lines(config, lines, table_cache=None, file_format_type: FileFormatT
  else:
  stream_to_sync[stream] = DbSync(config, o, table_cache, file_format_type)
 
+ if archive_load_files:
+ archive_load_files_data[stream] = {
+ 'tap': config.get('tap_id'),
+ }
+
+ # In case of incremental replication, track min/max of the replication key.
+ # Incremental replication is assumed if o['bookmark_properties'][0] is one of the columns.
+ incremental_key_column_name = stream_utils.get_incremental_key(o)
+ if incremental_key_column_name:
+ LOGGER.info("Using %s as incremental_key_column_name", incremental_key_column_name)
+ archive_load_files_data[stream].update(
+ column=incremental_key_column_name,
+ min=None,
+ max=None
+ )
+ else:
+ LOGGER.warning(
+ "archive_load_files is enabled, but no incremental_key_column_name was found. "
+ "Min/max values will not be added to metadata for stream %s.", stream
+ )
+
  stream_to_sync[stream].create_schema_if_not_exists()
  stream_to_sync[stream].sync_table()
 
@@ -281,7 +322,8 @@ def persist_lines(config, lines, table_cache=None, file_format_type: FileFormatT
  # then flush all buckets.
  if sum(row_count.values()) > 0:
  # flush all streams one last time, delete records if needed, reset counts and then emit current state
- flushed_state = flush_streams(records_to_load, row_count, stream_to_sync, config, state, flushed_state)
+ flushed_state = flush_streams(records_to_load, row_count, stream_to_sync, config, state, flushed_state,
+ archive_load_files_data)
 
  # emit latest state
  emit_state(copy.deepcopy(flushed_state))
@@ -295,6 +337,7 @@ def flush_streams(
  config,
  state,
  flushed_state,
+ archive_load_files_data,
  filter_streams=None):
  """
  Flushes all buckets and resets records count to 0 as well as empties records to load list
@@ -305,6 +348,7 @@ def flush_streams(
  :param state: dictionary containing the original state from tap
  :param flushed_state: dictionary containing updated states only when streams got flushed
  :param filter_streams: Keys of streams to flush from the streams dict. Default is every stream
+ :param archive_load_files_data: dictionary of dictionaries containing archive load files data
  :return: State dict with flushed positions
  """
  parallelism = config.get("parallelism", DEFAULT_PARALLELISM)
@@ -337,7 +381,8 @@ def flush_streams(
  db_sync=stream_to_sync[stream],
  no_compression=config.get('no_compression'),
  delete_rows=config.get('hard_delete'),
- temp_dir=config.get('temp_dir')
+ temp_dir=config.get('temp_dir'),
+ archive_load_files=copy.copy(archive_load_files_data.get(stream, None))
  ) for stream in streams_to_flush)
 
  # reset flushed stream records to empty to avoid flushing same records
@@ -358,16 +403,20 @@ def flush_streams(
  else:
  flushed_state = copy.deepcopy(state)
 
+ if stream in archive_load_files_data:
+ archive_load_files_data[stream]['min'] = None
+ archive_load_files_data[stream]['max'] = None
+
  # Return with state message with flushed positions
  return flushed_state
 
 
 def load_stream_batch(stream, records, row_count, db_sync, no_compression=False, delete_rows=False,
- temp_dir=None):
+ temp_dir=None, archive_load_files=None):
  """Load one batch of the stream into target table"""
  # Load into snowflake
  if row_count[stream] > 0:
- flush_records(stream, records, db_sync, temp_dir, no_compression)
+ flush_records(stream, records, db_sync, temp_dir, no_compression, archive_load_files)
 
  # Delete soft-deleted, flagged rows - where _sdc_deleted at is not null
  if delete_rows:
@@ -381,7 +430,8 @@ def flush_records(stream: str,
  records: List[Dict],
  db_sync: DbSync,
  temp_dir: str = None,
- no_compression: bool = False) -> 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
 
@@ -391,8 +441,9 @@ def flush_records(stream: str,
  column value
  row_count:
  db_sync: A DbSync object
- temp_dir: Directory where intermediate temporary files will be created. (Default: OS specificy temp directory)
+ 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
@@ -403,7 +454,7 @@ def flush_records(stream: str,
  compression=not no_compression,
  dest_dir=temp_dir,
  data_flattening_max_level=
-  db_sync.data_flattening_max_level)
+ db_sync.data_flattening_max_level)
 
  # Get file stats
  row_count = len(records)
@@ -413,8 +464,40 @@ def flush_records(stream: str,
  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 and from s3
+ # 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 = s3_key.split('/')[-1]
+ archive_folder = "archive/{}/{}".format(archive_tap, archive_table)
+ archive_key = "{}/{}".format(archive_folder, 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)
 
 

target_snowflake/db_sync.py

@@ -63,6 +63,11 @@ def validate_config(config):
  if not config_default_target_schema and not config_schema_mapping:
  errors.append("Neither 'default_target_schema' (string) nor 'schema_mapping' (object) keys set in config.")
 
+ # Check if archive load files option is using external stages
+ archive_load_files = config.get('archive_load_files', False)
+ if archive_load_files and not config.get('s3_bucket', None):
+ errors.append('Archive load files option can be used only with external s3 stages. Please define s3_bucket.')
+
  return errors
 
 
@@ -397,6 +402,11 @@ def delete_from_stage(self, stream, s3_key):
  self.logger.info('Deleting %s from stage', format(s3_key))
  self.upload_client.delete_object(stream, s3_key)
 
+ def copy_to_archive(self, s3_source_key, s3_archive_key, s3_archive_metadata):
+ """Copy file from snowflake stage to archive"""
+ self.logger.info('Copying %s to archive location %s', s3_source_key, s3_archive_key)
+ self.upload_client.copy_object(s3_source_key, s3_archive_key, s3_archive_metadata)
+
  def get_stage_name(self, stream):
  """Generate snowflake stage name"""
  stage = self.connection_config.get('stage', None)

target_snowflake/exceptions.py

@@ -23,3 +23,7 @@ class FileFormatNotFoundException(Exception):
 
 class InvalidFileFormatException(Exception):
  """Exception to raise when name file format is not compatible"""
+
+
+class UnexpectedMessageTypeException(Exception):
+ """Exception to raise when provided message doesn't match the expected type"""

target_snowflake/stream_utils.py

@@ -8,6 +8,7 @@
 from singer import get_logger
 
 from target_snowflake.exceptions import UnexpectedValueTypeException
+from target_snowflake.exceptions import UnexpectedMessageTypeException
 
 LOGGER = get_logger('target_snowflake')
 
@@ -114,3 +115,16 @@ def stream_name_to_dict(stream_name, separator='-'):
  'schema_name': schema_name,
  'table_name': table_name
  }
+
+
+def get_incremental_key(singer_msg: Dict):
+ """Derive incremental key from a Singer message dictionary"""
+ if singer_msg['type'] != "SCHEMA":
+ raise UnexpectedMessageTypeException("Expecting type SCHEMA, got {}".format(singer_msg['type']))
+
+ if 'bookmark_properties' in singer_msg and len(singer_msg['bookmark_properties']) > 0:
+ col = singer_msg['bookmark_properties'][0]
+ if col in singer_msg['schema']['properties']:
+ return col
+
+ return None

target_snowflake/upload_clients/base_upload_client.py

@@ -24,3 +24,9 @@ def delete_object(self, stream: str, key: str) -> None:
  """
  Delete object
  """
+
+ @abstractmethod
+ def copy_object(self, source_key: str, target_key: str, target_metadata: dict) -> None:
+ """
+ Copy object
+ """

target_snowflake/upload_clients/s3_upload_client.py

@@ -97,3 +97,14 @@ def delete_object(self, stream: str, key: str) -> None:
  self.logger.info('Deleting %s from external snowflake stage on S3', key)
  bucket = self.connection_config['s3_bucket']
  self.s3_client.delete_object(Bucket=bucket, Key=key)
+
+ def copy_object(self, source_key: str, target_key: str, target_metadata: dict) -> None:
+ """Copy object to another location on S3"""
+ self.logger.info('Copying %s to %s', source_key, target_key)
+ bucket = self.connection_config['s3_bucket']
+
+ copy_source = "{}/{}".format(bucket, source_key)
+
+ # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3.html#S3.Client.copy_object
+ self.s3_client.copy_object(CopySource=copy_source, Bucket=bucket, Key=target_key, Metadata=target_metadata,
+ MetadataDirective="REPLACE")

target_snowflake/upload_clients/snowflake_upload_client.py

@@ -38,3 +38,7 @@ def delete_object(self, stream: str, key: str) -> None:
 
  with self.dblink.open_connection() as connection:
  connection.cursor().execute(f"REMOVE '@{stage}/{key}'")
+
+ def copy_object(self, source_key: str, target_key: str, target_metadata: dict) -> None:
+ raise NotImplementedError(
+ "Copying objects is not supported with a Snowflake upload client.")