def test_constructor_yaml(mock_session, dbapi):
r = Redshift(profile=PROFILE, dbapi=dbapi, config_yaml="some_config.yml")
mock_session.assert_called_with(profile_name=PROFILE)
assert r.profile == PROFILE
assert r.kms_key == None
assert r.connection["host"] == "host"
assert r.connection["port"] == "port"
assert r.connection["database"] == "database"
assert r.connection["user"] == "user"
assert r.connection["password"] == "password"
def test_constructor(mock_session, credentials, dbapi):
r = Redshift(profile=PROFILE, dbapi=dbapi, **credentials)
mock_session.assert_called_with(profile_name=PROFILE)
assert r.profile == PROFILE
assert r.kms_key == None
assert r.connection["host"] == "host"
assert r.connection["port"] == "port"
assert r.connection["database"] == "database"
assert r.connection["user"] == "user"
assert r.connection["password"] == "password"
def test_redshift_connect(cls,
mock_session,
credentials=credentials(),
dbapi=DBAPIS):
with mock.patch(dbapi.__name__ + ".connect") as mock_connect:
r = Redshift(profile=PROFILE, dbapi=dbapi, **credentials)
r.connect()
if dbapi.__name__ == "pg8000":
mock_connect.assert_called_with(
host="host",
user="******",
port="port",
password="******",
database="database",
ssl=True,
)
else:
mock_connect.assert_called_with(
host="host",
user="******",
port="port",
password="******",
database="database",
sslmode="require",
)
r.conn.cursor.assert_called_with()
# side effect exception
mock_connect.side_effect = Exception("Connect Exception")
with pytest.raises(DBError):
r.connect()
def test_load_and_copy(
cls,
mock_split_file,
mock_compress_file_list,
mock_session,
mock_s3_delete,
mock_s3_upload,
mock_rs_copy,
mock_remove,
credentials=credentials(),
dbapi=DBAPIS,
):
def reset_mocks():
mock_split_file.reset_mock()
mock_compress_file_list.reset_mock()
mock_s3_upload.reset_mock()
mock_s3_delete.reset_mock()
mock_rs_copy.reset_mock()
mock_remove.reset_mock()
with mock.patch(dbapi.__name__ + ".connect") as mock_connect:
r = Redshift(dbapi=dbapi, **credentials)
r.connect()
expected_calls_no_folder = [
mock.call("/path/local_file.0", "s3_bucket", "local_file.0"),
mock.call("/path/local_file.1", "s3_bucket", "local_file.1"),
mock.call("/path/local_file.2", "s3_bucket", "local_file.2"),
]
expected_calls_no_folder_gzip = [
mock.call("/path/local_file.0.gz", "s3_bucket",
"local_file.0.gz"),
mock.call("/path/local_file.1.gz", "s3_bucket",
"local_file.1.gz"),
mock.call("/path/local_file.2.gz", "s3_bucket",
"local_file.2.gz"),
]
expected_calls_folder = [
mock.call("/path/local_file.0", "s3_bucket",
"test/local_file.0"),
mock.call("/path/local_file.1", "s3_bucket",
"test/local_file.1"),
mock.call("/path/local_file.2", "s3_bucket",
"test/local_file.2"),
]
expected_calls_folder_gzip = [
mock.call("/path/local_file.0.gz", "s3_bucket",
"test/local_file.0.gz"),
mock.call("/path/local_file.1.gz", "s3_bucket",
"test/local_file.1.gz"),
mock.call("/path/local_file.2.gz", "s3_bucket",
"test/local_file.2.gz"),
]
mock_split_file.return_value = ["/path/local_file.txt"]
mock_compress_file_list.return_value = ["/path/local_file.txt.gz"]
r.load_and_copy("/path/local_file.txt",
"s3_bucket",
"table_name",
delim="|")
# assert
assert mock_split_file.called
mock_compress_file_list.assert_called_with(
["/path/local_file.txt"])
# mock_remove.assert_called_with("/path/local_file.txt")
mock_s3_upload.assert_called_with("/path/local_file.txt.gz",
"s3_bucket", "local_file.txt.gz")
mock_rs_copy.assert_called_with("table_name",
"s3://s3_bucket/local_file",
"|",
copy_options=["GZIP"])
assert not mock_s3_delete.called, "Only delete when explicit"
reset_mocks()
mock_split_file.return_value = [
"/path/local_file.0",
"/path/local_file.1",
"/path/local_file.2",
]
mock_compress_file_list.return_value = [
"/path/local_file.0.gz",
"/path/local_file.1.gz",
"/path/local_file.2.gz",
]
r.load_and_copy(
"/path/local_file",
"s3_bucket",
"table_name",
delim="|",
copy_options=["SOME OPTION"],
splits=3,
delete_s3_after=True,
)
# assert
mock_split_file.assert_called_with("/path/local_file",
"/path/local_file",
splits=3)
mock_compress_file_list.assert_called_with([
"/path/local_file.0", "/path/local_file.1",
"/path/local_file.2"
])
# mock_remove.assert_called_with("/path/local_file.2")
mock_s3_upload.assert_has_calls(expected_calls_no_folder_gzip)
mock_rs_copy.assert_called_with(
"table_name",
"s3://s3_bucket/local_file",
"|",
copy_options=["SOME OPTION", "GZIP"])
assert mock_s3_delete.called_with("s3_bucket", "local_file.0.gz")
assert mock_s3_delete.called_with("s3_bucket", "local_file.1.gz")
assert mock_s3_delete.called_with("s3_bucket", "local_file.2.gz")
reset_mocks()
mock_split_file.return_value = ["/path/local_file"]
mock_compress_file_list.return_value = ["/path/local_file.gz"]
r.load_and_copy(
"/path/local_file",
"s3_bucket",
"table_name",
delim=",",
copy_options=["SOME OPTION"],
compress=False,
)
# assert
assert mock_split_file.called
assert not mock_compress_file_list.called
# assert not mock_remove.called
mock_s3_upload.assert_called_with("/path/local_file", "s3_bucket",
"local_file")
mock_rs_copy.assert_called_with("table_name",
"s3://s3_bucket/local_file",
",",
copy_options=["SOME OPTION"])
assert not mock_s3_delete.called, "Only delete when explicit"
reset_mocks()
mock_split_file.return_value = [
"/path/local_file.0",
"/path/local_file.1",
"/path/local_file.2",
]
r.load_and_copy(
"/path/local_file",
"s3_bucket",
"table_name",
delim="|",
copy_options=["SOME OPTION"],
splits=3,
compress=False,
)
# assert
mock_split_file.assert_called_with("/path/local_file",
"/path/local_file",
splits=3)
assert not mock_compress_file_list.called
# assert not mock_remove.called
mock_s3_upload.assert_has_calls(expected_calls_no_folder)
mock_rs_copy.assert_called_with("table_name",
"s3://s3_bucket/local_file",
"|",
copy_options=["SOME OPTION"])
assert not mock_s3_delete.called
# with a s3_folder included and no splits
reset_mocks()
mock_split_file.return_value = ["/path/local_file.txt"]
r.load_and_copy(
"/path/local_file.txt",
"s3_bucket",
"table_name",
delim="|",
copy_options=["SOME OPTION"],
compress=False,
s3_folder="test",
)
# assert
assert mock_split_file.called
assert not mock_compress_file_list.called
# assert not mock_remove.called
mock_s3_upload.assert_called_with("/path/local_file.txt",
"s3_bucket",
"test/local_file.txt")
mock_rs_copy.assert_called_with("table_name",
"s3://s3_bucket/test/local_file",
"|",
copy_options=["SOME OPTION"])
assert not mock_s3_delete.called
# with a s3_folder included and splits
reset_mocks()
mock_split_file.return_value = [
"/path/local_file.0",
"/path/local_file.1",
"/path/local_file.2",
]
r.load_and_copy(
"/path/local_file",
"s3_bucket",
"table_name",
delim="|",
copy_options=["SOME OPTION"],
splits=3,
compress=False,
s3_folder="test",
delete_s3_after=True,
)
# assert
mock_split_file.assert_called_with("/path/local_file",
"/path/local_file",
splits=3)
assert not mock_compress_file_list.called
# assert not mock_remove.called
mock_s3_upload.assert_has_calls(expected_calls_folder)
mock_rs_copy.assert_called_with("table_name",
"s3://s3_bucket/test/local_file",
"|",
copy_options=["SOME OPTION"])
assert mock_s3_delete.called_with("s3_bucket", "test/local_file.0")
assert mock_s3_delete.called_with("s3_bucket", "test/local_file.1")
assert mock_s3_delete.called_with("s3_bucket", "test/local_file.2")
# with a s3_folder included , splits, and gzip
reset_mocks()
mock_split_file.return_value = [
"/path/local_file.0",
"/path/local_file.1",
"/path/local_file.2",
]
mock_compress_file_list.return_value = [
"/path/local_file.0.gz",
"/path/local_file.1.gz",
"/path/local_file.2.gz",
]
r.load_and_copy(
"/path/local_file",
"s3_bucket",
"table_name",
delim="|",
copy_options=["SOME OPTION"],
splits=3,
s3_folder="test",
)
# assert
mock_split_file.assert_called_with("/path/local_file",
"/path/local_file",
splits=3)
assert mock_compress_file_list.called
# assert mock_remove.called
mock_s3_upload.assert_has_calls(expected_calls_folder_gzip)
mock_rs_copy.assert_called_with(
"table_name",
"s3://s3_bucket/test/local_file",
"|",
copy_options=["SOME OPTION", "GZIP"],
)
assert not mock_s3_delete.called
def test_load_and_copy_split_and_header(
mock_split_file,
mock_compress_file_list,
mock_session,
mock_s3_delete,
mock_s3_upload,
mock_rs_copy,
mock_remove,
credentials,
dbapi,
):
def reset_mocks():
mock_split_file.reset_mock()
mock_compress_file_list.reset_mock()
mock_s3_upload.reset_mock()
mock_s3_delete.reset_mock()
mock_rs_copy.reset_mock()
mock_remove.reset_mock()
with mock.patch(dbapi.__name__ + ".connect") as mock_connect:
r = Redshift(dbapi=dbapi, **credentials)
r.connect()
expected_calls_no_folder = [
mock.call("/path/local_file.0", "s3_bucket", "local_file.0"),
mock.call("/path/local_file.1", "s3_bucket", "local_file.1"),
mock.call("/path/local_file.2", "s3_bucket", "local_file.2"),
]
expected_calls_no_folder_gzip = [
mock.call("/path/local_file.0.gz", "s3_bucket", "local_file.0.gz"),
mock.call("/path/local_file.1.gz", "s3_bucket", "local_file.1.gz"),
mock.call("/path/local_file.2.gz", "s3_bucket", "local_file.2.gz"),
]
expected_calls_folder = [
mock.call("/path/local_file.0", "s3_bucket", "test/local_file.0"),
mock.call("/path/local_file.1", "s3_bucket", "test/local_file.1"),
mock.call("/path/local_file.2", "s3_bucket", "test/local_file.2"),
]
expected_calls_folder_gzip = [
mock.call("/path/local_file.0.gz", "s3_bucket",
"test/local_file.0.gz"),
mock.call("/path/local_file.1.gz", "s3_bucket",
"test/local_file.1.gz"),
mock.call("/path/local_file.2.gz", "s3_bucket",
"test/local_file.2.gz"),
]
mock_split_file.return_value = ["/path/local_file.txt"]
mock_compress_file_list.return_value = ["/path/local_file.txt.gz"]
#### neither ignore or split only
r.load_and_copy("/path/local_file.txt",
"s3_bucket",
"table_name",
delim="|")
# assert
assert mock_split_file.called
mock_compress_file_list.assert_called_with(["/path/local_file.txt"])
# mock_remove.assert_called_with("/path/local_file.txt")
mock_s3_upload.assert_called_with("/path/local_file.txt.gz",
"s3_bucket", "local_file.txt.gz")
mock_rs_copy.assert_called_with("table_name",
"s3://s3_bucket/local_file",
"|",
copy_options=["GZIP"])
assert not mock_s3_delete.called, "Only delete when explicit"
#### ignore only
reset_mocks()
r.load_and_copy(
"/path/local_file.txt",
"s3_bucket",
"table_name",
delim="|",
copy_options=["IGNOREHEADER as 1"],
)
# assert
assert mock_split_file.called
mock_compress_file_list.assert_called_with(["/path/local_file.txt"])
# mock_remove.assert_called_with("/path/local_file.txt")
mock_s3_upload.assert_called_with("/path/local_file.txt.gz",
"s3_bucket", "local_file.txt.gz")
mock_rs_copy.assert_called_with(
"table_name",
"s3://s3_bucket/local_file",
"|",
copy_options=["IGNOREHEADER as 1", "GZIP"],
)
assert not mock_s3_delete.called, "Only delete when explicit"
#### split only
reset_mocks()
mock_split_file.return_value = [
"/path/local_file.0",
"/path/local_file.1",
"/path/local_file.2",
]
mock_compress_file_list.return_value = [
"/path/local_file.0.gz",
"/path/local_file.1.gz",
"/path/local_file.2.gz",
]
r.load_and_copy(
"/path/local_file",
"s3_bucket",
"table_name",
delim="|",
splits=3,
delete_s3_after=True,
)
# assert
mock_split_file.assert_called_with("/path/local_file",
"/path/local_file",
splits=3,
ignore_header=0)
mock_compress_file_list.assert_called_with(
["/path/local_file.0", "/path/local_file.1", "/path/local_file.2"])
mock_s3_upload.assert_has_calls(expected_calls_no_folder_gzip)
mock_rs_copy.assert_called_with("table_name",
"s3://s3_bucket/local_file",
"|",
copy_options=["GZIP"])
assert mock_s3_delete.called_with("s3_bucket", "local_file.0.gz")
assert mock_s3_delete.called_with("s3_bucket", "local_file.1.gz")
assert mock_s3_delete.called_with("s3_bucket", "local_file.2.gz")
#### split and ignore
reset_mocks()
mock_split_file.return_value = [
"/path/local_file.0",
"/path/local_file.1",
"/path/local_file.2",
]
mock_compress_file_list.return_value = [
"/path/local_file.0.gz",
"/path/local_file.1.gz",
"/path/local_file.2.gz",
]
r.load_and_copy(
"/path/local_file",
"s3_bucket",
"table_name",
delim="|",
copy_options=["IGNOREHEADER as 1"],
splits=3,
delete_s3_after=True,
)
# assert
mock_split_file.assert_called_with("/path/local_file",
"/path/local_file",
splits=3,
ignore_header=1)
mock_compress_file_list.assert_called_with(
["/path/local_file.0", "/path/local_file.1", "/path/local_file.2"])
# mock_remove.assert_called_with("/path/local_file.2")
mock_s3_upload.assert_has_calls(expected_calls_no_folder_gzip)
mock_rs_copy.assert_called_with("table_name",
"s3://s3_bucket/local_file",
"|",
copy_options=["GZIP"])
assert mock_s3_delete.called_with("s3_bucket", "local_file.0.gz")
assert mock_s3_delete.called_with("s3_bucket", "local_file.1.gz")
assert mock_s3_delete.called_with("s3_bucket", "local_file.2.gz")
def test_constructor_credential_error(credentials, dbapi, caplog):
r = Redshift(dbapi=dbapi, **credentials)
assert "S3 credentials were not found. S3 functionality is disabled" in caplog.text