def test_its_created_successfully(self):
operator = PostgresToS3Transfer(
task_id='task_id',
postgres_conn_id='postgres_conn_id',
s3_conn_id='s3_conn_id',
s3_url='s3://bucket/key'
)
assert operator
assert operator.task_id == 'task_id'
def test_execute_streams_url_data_to_s3(self, get_connection_mock,
boto3_mock, popen_mock):
operator = PostgresToS3Transfer(task_id='task_id',
postgres_conn_id='postgres_conn_id',
s3_conn_id='s3_conn_id',
s3_url='s3://bucket/key')
operator.execute(None)
boto3_mock().Bucket.assert_called_with('bucket')
boto3_mock().Bucket().upload_fileobj.assert_called_with(
popen_mock().stdout.__enter__(
), # Needs __enter__() because it's called in a context manager
'key')
def test_execute_calls_pg_dump_correctly(self, get_connection_mock,
boto3_mock, popen_mock):
operator = PostgresToS3Transfer(task_id='task_id',
postgres_conn_id='postgres_conn_id',
s3_conn_id='s3_conn_id',
s3_url='s3://bucket/key')
operator.execute(None)
expected_command = [
'pg_dump',
'-Fc',
helpers.get_postgres_uri(operator.postgres_conn_id),
]
popen_mock.assert_called_with(expected_command, stdout=subprocess.PIPE)
def test_execute_dumps_only_whitelisted_tables(self, get_connection_mock,
boto3_mock, popen_mock):
tables = [
'users',
'log',
]
operator = PostgresToS3Transfer(task_id='task_id',
postgres_conn_id='postgres_conn_id',
tables=tables,
s3_conn_id='s3_conn_id',
s3_url='s3://bucket/key')
operator.execute(None)
popen_command = popen_mock.call_args[0][0]
# Ignore executable and the Postgres URI, as the params need to be
# between these two
pg_dump_params_without_uri = popen_command[1:-1]
for table in tables:
assert '--table={}'.format(table) in pg_dump_params_without_uri