AP1011 - Custom archive bucket/folder

README.md

@@ -176,7 +176,9 @@ 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`. 
+| archive_load_files | Boolean | | (Default: False) When enabled, the files loaded to Snowflake will also be stored in `archive_load_files_s3_bucket` under the key `/{archive_load_files_s3_prefix}/{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`. 
+| archive_load_files_s3_prefix | String | | (Default: "archive") When `archive_load_files` is enabled, the archived files will be placed in the archive S3 bucket under this prefix.
+| archive_load_files_s3_bucket | String | | (Default: Value of `s3_bucket`) When `archive_load_files` is enabled, the archived files will be placed in this bucket.
 
 ### To run tests:
 

target_snowflake/__init__.py

@@ -491,9 +491,8 @@ def flush_records(stream: str,
  })
 
  # 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)
+ 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)
 

target_snowflake/db_sync.py

@@ -403,9 +403,35 @@ def delete_from_stage(self, stream, 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)
+ """
+ Copy file from snowflake stage to archive.
+
+ s3_source_key: The s3 key to copy, assumed to exist in the bucket configured as 's3_bucket'
+
+ s3_archive_key: The key to use in archive destination. This will be prefixed with the config value
+ 'archive_load_files_s3_prefix'. If none is specified, 'archive' will be used as the prefix.
+
+ As destination bucket, the config value 'archive_load_files_s3_bucket' will be used. If none is
+ specified, the bucket configured as 's3_bucket' will be used.
+
+ s3_archive_metadata: This dict will be used as the S3 metadata in the file in archive destination. Metadata in
+ the source file will be replaced.
+
+ """
+ source_bucket = self.connection_config.get('s3_bucket')
+
+ # Get archive s3_bucket from config, or use same bucket if not specified
+ archive_bucket = self.connection_config.get('archive_load_files_s3_bucket', source_bucket)
+
+ # Determine prefix to use in archive s3 bucket
+ default_archive_prefix = 'archive'
+ archive_prefix = self.connection_config.get('archive_load_files_s3_prefix', default_archive_prefix)
+ prefixed_archive_key = '{}/{}'.format(archive_prefix, s3_archive_key)
+
+ copy_source = '{}/{}'.format(source_bucket, s3_source_key)
+
+ self.logger.info('Copying %s to archive location %s', copy_source, prefixed_archive_key)
+ self.upload_client.copy_object(copy_source, archive_bucket, prefixed_archive_key, s3_archive_metadata)
 
  def get_stage_name(self, stream):
  """Generate snowflake stage name"""

target_snowflake/upload_clients/base_upload_client.py

@@ -26,7 +26,7 @@ def delete_object(self, stream: str, key: str) -> None:
  """
 
  @abstractmethod
- def copy_object(self, source_key: str, target_key: str, target_metadata: dict) -> None:
+ def copy_object(self, copy_source: str, target_bucket: str, target_key: str, target_metadata: dict) -> None:
  """
  Copy object
  """

target_snowflake/upload_clients/s3_upload_client.py

@@ -98,13 +98,9 @@ def delete_object(self, stream: str, key: str) -> None:
  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:
+ def copy_object(self, copy_source: str, target_bucket: 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)
-
+ self.logger.info('Copying %s to %s/%s', copy_source, target_bucket, target_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")
+ self.s3_client.copy_object(CopySource=copy_source, Bucket=target_bucket, Key=target_key,
+ Metadata=target_metadata, MetadataDirective="REPLACE")

target_snowflake/upload_clients/snowflake_upload_client.py

@@ -39,6 +39,6 @@ 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:
+ def copy_object(self, copy_source: str, target_bucket: str, target_key: str, target_metadata: dict) -> None:
  raise NotImplementedError(
  "Copying objects is not supported with a Snowflake upload client.")

tests/integration/test_target_snowflake.py

@@ -1180,14 +1180,15 @@ def test_parquet(self):
  def test_archive_load_files(self):
  """Test if load file is copied to archive folder"""
  self.config['archive_load_files'] = True
+ self.config['archive_load_files_s3_prefix'] = 'archive_folder'
  self.config['tap_id'] = 'test_tap_id'
  self.config['client_side_encryption_master_key'] = ''
 
  s3_bucket = self.config['s3_bucket']
 
  # Delete any dangling files from archive
  files_in_s3_archive = self.s3_client.list_objects(
- Bucket=s3_bucket, Prefix="archive/test_tap_id/").get('Contents', [])
+ Bucket=s3_bucket, Prefix="archive_folder/test_tap_id/").get('Contents', [])
  for file_in_archive in files_in_s3_archive:
  key = file_in_archive["Key"]
  self.s3_client.delete_object(Bucket=s3_bucket, Key=key)
@@ -1196,7 +1197,7 @@ def test_archive_load_files(self):
  self.persist_lines_with_cache(tap_lines)
 
  # Verify expected file metadata in S3
- files_in_s3_archive = self.s3_client.list_objects(Bucket=s3_bucket, Prefix="archive/test_tap_id/").get(
+ files_in_s3_archive = self.s3_client.list_objects(Bucket=s3_bucket, Prefix="archive_folder/test_tap_id/").get(
  'Contents')
  self.assertIsNotNone(files_in_s3_archive)
  self.assertEqual(1, len(files_in_s3_archive))
@@ -1229,6 +1230,3 @@ def test_archive_load_files(self):
 4,"xyz4","not-formatted-time-4"
 5,"xyz5","not-formatted-time-5"
 ''')
-
- # Clean up
- self.s3_client.delete_object(Bucket=s3_bucket, Key=archived_file_key)

tests/unit/test_db_sync.py

@@ -224,6 +224,43 @@ def test_parallelism(self, query_patch):
  self.assertEqual(db_sync.DbSync({**minimal_config,
  **table_stage_with_parallel}).connection_config['parallelism'], 1)
 
+ @patch('target_snowflake.upload_clients.s3_upload_client.S3UploadClient.copy_object')
+ @patch('target_snowflake.db_sync.DbSync.query')
+ def test_copy_to_archive(self, query_patch, copy_object_patch):
+ query_patch.return_value = [{'type': 'CSV'}]
+ minimal_config = {
+ 'account': "dummy-value",
+ 'dbname': "dummy-value",
+ 'user': "dummy-value",
+ 'password': "dummy-value",
+ 'warehouse': "dummy-value",
+ 'default_target_schema': "dummy-value",
+ 'file_format': "dummy-value",
+ 's3_bucket': 'dummy-bucket',
+ 'stage': 'dummy_schema.dummy_stage'
+ }
+
+ # Assert default values (same bucket, 'archive' as the archive prefix)
+ s3_config = {}
+ dbsync = db_sync.DbSync({**minimal_config, **s3_config})
+ dbsync.copy_to_archive('source/file', 'tap/schema/file', {'meta': "data"})
+
+ self.assertEqual(copy_object_patch.call_args[0][0], 'dummy-bucket/source/file')
+ self.assertEqual(copy_object_patch.call_args[0][1], 'dummy-bucket')
+ self.assertEqual(copy_object_patch.call_args[0][2], 'archive/tap/schema/file')
+
+ # Assert custom archive bucket and prefix
+ s3_config = {
+ 'archive_load_files_s3_bucket': "custom-bucket",
+ 'archive_load_files_s3_prefix': "custom-prefix"
+ }
+ dbsync = db_sync.DbSync({**minimal_config, **s3_config})
+ dbsync.copy_to_archive('source/file', 'tap/schema/file', {'meta': "data"})
+
+ self.assertEqual(copy_object_patch.call_args[0][0], 'dummy-bucket/source/file')
+ self.assertEqual(copy_object_patch.call_args[0][1], 'custom-bucket')
+ self.assertEqual(copy_object_patch.call_args[0][2], 'custom-prefix/tap/schema/file')
+
  def test_safe_column_name(self):
  self.assertEqual(db_sync.safe_column_name("columnname"), '"COLUMNNAME"')
  self.assertEqual(db_sync.safe_column_name("columnName"), '"COLUMNNAME"')

tests/unit/test_target_snowflake.py

@@ -100,13 +100,13 @@ def test_archive_load_files_incremental_replication(self, os_remove_mock, dbSync
  instance = dbSync_mock.return_value
  instance.create_schema_if_not_exists.return_value = None
  instance.sync_table.return_value = None
- instance.put_to_stage.return_value = 'some-s3-bucket/some-s3-folder/some-name_date_batch_hash.csg.gz'
+ instance.put_to_stage.return_value = 'some-s3-folder/some-name_date_batch_hash.csg.gz'
 
  target_snowflake.persist_lines(self.config, lines)
 
  copy_to_archive_args = instance.copy_to_archive.call_args[0]
- assert copy_to_archive_args[0] == 'some-s3-bucket/some-s3-folder/some-name_date_batch_hash.csg.gz'
- assert copy_to_archive_args[1] == 'archive/test_tap_id/test_simple_table/some-name_date_batch_hash.csg.gz'
+ assert copy_to_archive_args[0] == 'some-s3-folder/some-name_date_batch_hash.csg.gz'
+ assert copy_to_archive_args[1] == 'test_tap_id/test_simple_table/some-name_date_batch_hash.csg.gz'
  assert copy_to_archive_args[2] == {
  'tap': 'test_tap_id',
  'schema': 'tap_mysql_test',
@@ -122,21 +122,20 @@ def test_archive_load_files_incremental_replication(self, os_remove_mock, dbSync
  def test_archive_load_files_log_based_replication(self, os_remove_mock, dbSync_mock):
  self.config['tap_id'] = 'test_tap_id'
  self.config['archive_load_files'] = True
- self.config['s3_bucket'] = 'dummy_bucket'
 
  with open(f'{os.path.dirname(__file__)}/resources/logical-streams.json', 'r') as f:
  lines = f.readlines()
 
  instance = dbSync_mock.return_value
  instance.create_schema_if_not_exists.return_value = None
  instance.sync_table.return_value = None
- instance.put_to_stage.return_value = 'some-s3-bucket/some-s3-folder/some-name_date_batch_hash.csg.gz'
+ instance.put_to_stage.return_value = 'some-s3-folder/some-name_date_batch_hash.csg.gz'
 
  target_snowflake.persist_lines(self.config, lines)
 
  copy_to_archive_args = instance.copy_to_archive.call_args[0]
- assert copy_to_archive_args[0] == 'some-s3-bucket/some-s3-folder/some-name_date_batch_hash.csg.gz'
- assert copy_to_archive_args[1] == 'archive/test_tap_id/logical1_table2/some-name_date_batch_hash.csg.gz'
+ assert copy_to_archive_args[0] == 'some-s3-folder/some-name_date_batch_hash.csg.gz'
+ assert copy_to_archive_args[1] == 'test_tap_id/logical1_table2/some-name_date_batch_hash.csg.gz'
  assert copy_to_archive_args[2] == {
  'tap': 'test_tap_id',
  'schema': 'logical1',