def test_export_snowflake_table_to_s3_with_binary_format(
mock_sf_connection): # noqa: F811
mock_cursor = mock_sf_connection.cursor()
with Flow("test") as f:
snowflake.export_snowflake_table_to_s3(
sf_credentials={},
sf_database="test_database",
sf_schema="test_schema",
sf_table="test_table",
sf_role="test_role",
sf_warehouse="test_warehouse",
sf_storage_integration="test_storage_integration",
s3_path="s3://edx-test/test/",
overwrite=False,
enclosed_by='"',
escape_unenclosed_field='\\\\',
null_marker='NULL',
binary_format='UTF8',
)
state = f.run()
assert state.is_successful()
mock_cursor.execute.assert_has_calls([
mock.call(
"""\n COPY INTO 's3://edx-test/test/test_database-test_schema-test_table/'\n FROM test_database.test_schema.test_table\n STORAGE_INTEGRATION = test_storage_integration\n FILE_FORMAT = ( TYPE = CSV EMPTY_FIELD_AS_NULL = FALSE\n FIELD_DELIMITER = ',' FIELD_OPTIONALLY_ENCLOSED_BY = '"'\n ESCAPE_UNENCLOSED_FIELD = '\\\\'\n NULL_IF = ( 'NULL' )\n BINARY_FORMAT = UTF8\n COMPRESSION = NONE\n )\n OVERWRITE=False\n SINGLE=False\n DETAILED_OUTPUT = TRUE\n MAX_FILE_SIZE = 104857600\n """
), # noqa
])
def test_export_snowflake_table_to_s3_no_null_if(
mock_sf_connection): # noqa: F811
mock_cursor = mock_sf_connection.cursor()
with mock.patch('edx_prefectutils.s3.delete_s3_directory.run'):
with Flow("test") as f:
snowflake.export_snowflake_table_to_s3(
sf_credentials={},
sf_database="test_database",
sf_schema="test_schema",
sf_table="test_table",
sf_role="test_role",
sf_warehouse="test_warehouse",
sf_storage_integration="test_storage_integration",
s3_path="s3://edx-test/test/",
overwrite=True,
enclosed_by='NONE',
escape_unenclosed_field='\\\\',
)
state = f.run()
assert state.is_successful()
mock_cursor.execute.assert_has_calls([
mock.call(
"\n COPY INTO 's3://edx-test/test/test_database-test_schema-test_table/'\n FROM test_database.test_schema.test_table\n STORAGE_INTEGRATION = test_storage_integration\n FILE_FORMAT = ( TYPE = CSV EMPTY_FIELD_AS_NULL = FALSE\n FIELD_DELIMITER = ',' FIELD_OPTIONALLY_ENCLOSED_BY = NONE\n ESCAPE_UNENCLOSED_FIELD = '\\\\'\n \n \n COMPRESSION = NONE\n )\n OVERWRITE=True\n SINGLE=False\n DETAILED_OUTPUT = TRUE\n MAX_FILE_SIZE = 104857600\n "
), # noqa
])
def test_export_snowflake_table_to_s3_with_manifest(
mock_sf_connection): # noqa: F811
mock_cursor = mock_sf_connection.cursor()
mock_fetchall = mock.Mock()
s3_files = ['data_0_0_0.csv', 'data_0_0_1.csv']
mock_fetchall.return_value = [[file] for file in s3_files]
mock_cursor.fetchall = mock_fetchall
with mock.patch('prefect.tasks.aws.s3.S3Upload.run') as mock_s3_upload:
with Flow("test") as f:
snowflake.export_snowflake_table_to_s3(
sf_credentials={},
sf_database="test_database",
sf_schema="test_schema",
sf_table="test_table",
sf_role="test_role",
sf_warehouse="test_warehouse",
sf_storage_integration="test_storage_integration",
s3_path="s3://edx-test/test/",
overwrite=False,
generate_manifest=True,
)
state = f.run()
assert state.is_successful()
expected_manifest_content = {
"entries": [
{
"url":
"s3://edx-test/test/test_database-test_schema-test_table/"
+ s3_file,
"mandatory":
True
} for s3_file in s3_files # noqa
]
}
mock_s3_upload.assert_called_once_with(
json.dumps(expected_manifest_content),
key="test/test_database-test_schema-test_table/manifest.json")