def test_upload_to_internal(mock_session, sf_credentials):
with mock.patch("snowflake.connector.connect") as mock_connect:
with Snowflake(profile=PROFILE, dbapi=DBAPIS, **sf_credentials) as sf:
sf.upload_to_internal("/some/file", "@~/internal")
sf.conn.cursor.return_value.execute.assert_called_with(
"PUT 'file:///some/file' @~/internal PARALLEL=4 AUTO_COMPRESS=True OVERWRITE=True",
(),
)
sf.upload_to_internal("/some/file", "@~/internal", parallel=99, auto_compress=False)
sf.conn.cursor.return_value.execute.assert_called_with(
"PUT 'file:///some/file' @~/internal PARALLEL=99 AUTO_COMPRESS=False OVERWRITE=True",
(),
)
sf.upload_to_internal("/some/file", "@~/internal", overwrite=False)
sf.conn.cursor.return_value.execute.assert_called_with(
"PUT 'file:///some/file' @~/internal PARALLEL=4 AUTO_COMPRESS=True OVERWRITE=False",
(),
)
# exception
sf.conn.cursor.return_value.execute.side_effect = Exception("PUT Exception")
with pytest.raises(DBError):
sf.upload_to_internal("/some/file", "@~/internal")
def test_download_from_internal_windows(mock_session, sf_credentials):
with mock.patch("snowflake.connector.connect") as mock_connect:
with Snowflake(profile=PROFILE, dbapi=DBAPIS, **sf_credentials) as sf:
sf.download_from_internal("@~/internal", r"C:\some\file")
sf.conn.cursor.return_value.execute.assert_called_with(
"GET @~/internal 'file://C:/some/file' PARALLEL=10", ())
def test_unload(
mock_session,
file_type,
format_options,
header,
copy_options,
expected,
sf_credentials,
):
with mock.patch("snowflake.connector.connect") as mock_connect:
with Snowflake(profile=PROFILE, dbapi=DBAPIS, **sf_credentials) as sf:
sf.unload(
"@~/stage",
"table_name",
file_type=file_type,
format_options=format_options,
header=header,
copy_options=copy_options,
)
sf.conn.cursor.return_value.execute.assert_called_with(
"COPY INTO @~/stage FROM table_name FILE_FORMAT = {0}".format(
expected),
(),
)
def test_upload_to_internal_windows(mock_session, sf_credentials):
with mock.patch("snowflake.connector.connect") as mock_connect:
with Snowflake(profile=PROFILE, dbapi=DBAPIS, **sf_credentials) as sf:
sf.upload_to_internal(r"C:\some\file", "@~/internal")
sf.conn.cursor.return_value.execute.assert_called_with(
"PUT 'file://*****:*****@~/internal PARALLEL=4 AUTO_COMPRESS=True",
())
def test_constructor(input_kms_key, profile, mock_session, sf_credentials):
sf = Snowflake(profile=profile, kms_key=input_kms_key, dbapi=DBAPIS, **sf_credentials)
mock_session.assert_called_with(profile_name=profile)
assert sf.profile == profile
assert sf.kms_key == input_kms_key
assert sf.connection["account"] == "account"
assert sf.connection["warehouse"] == "warehouse"
assert sf.connection["database"] == "database"
assert sf.connection["schema"] == "schema"
assert sf.connection["user"] == "user"
assert sf.connection["password"] == "password"
def test_insert_dataframe_to_table(mock_session, sf_credentials):
import pandas as pd
test_df = pd.read_csv(os.path.join(CURR_DIR, "data", "mock_dataframe.txt"), sep=",")
with mock.patch("snowflake.connector.connect") as mock_connect:
with Snowflake(profile=PROFILE, dbapi=DBAPIS, **sf_credentials) as sf:
sf.insert_dataframe_to_table(test_df, "database.schema.test")
sf.conn.cursor.return_value.executemany.assert_called_with(
"INSERT INTO database.schema.test (a,b,c) VALUES (%s,%s,%s)",
[("1", "x", "2011-01-01"), ("2", "y", "2001-04-02")],
)
sf.insert_dataframe_to_table(test_df, "database.schema.test", create=True)
sf.conn.cursor.return_value.execute.assert_any_call(
"CREATE TABLE database.schema.test (a int,b varchar,c date)", ()
)
sf.conn.cursor.return_value.executemany.assert_called_with(
"INSERT INTO database.schema.test (a,b,c) VALUES (%s,%s,%s)",
[("1", "x", "2011-01-01"), ("2", "y", "2001-04-02")],
)
sf.insert_dataframe_to_table(test_df, "database.schema.test", columns=["a", "b"])
sf.conn.cursor.return_value.executemany.assert_called_with(
"INSERT INTO database.schema.test (a,b) VALUES (%s,%s)", [("1", "x"), ("2", "y")]
)
sf.insert_dataframe_to_table(
test_df,
"database.schema.test",
create=True,
metadata=OrderedDict([("col1", "int"), ("col2", "varchar"), ("col3", "date")]),
)
sf.conn.cursor.return_value.execute.assert_any_call(
"CREATE TABLE database.schema.test (col1 int,col2 varchar,col3 date)", ()
)
sf.conn.cursor.return_value.executemany.assert_called_with(
"INSERT INTO database.schema.test (col1,col2,col3) VALUES (%s,%s,%s)",
[("1", "x", "2011-01-01"), ("2", "y", "2001-04-02")],
)
sf.insert_dataframe_to_table(
test_df,
"database.schema.test",
create=False,
metadata=OrderedDict([("col1", "int"), ("col2", "varchar"), ("col3", "date")]),
)
# mock_session.warn.assert_called_with('Metadata will not be used because create is set to False.')
sf.conn.cursor.return_value.executemany.assert_called_with(
"INSERT INTO database.schema.test (a,b,c) VALUES (%s,%s,%s)",
[("1", "x", "2011-01-01"), ("2", "y", "2001-04-02")],
)
def test_to_pandas(mock_session, sf_credentials):
import pandas as pd
test_df = pd.read_csv(os.path.join(CURR_DIR, "data", "mock_dataframe.txt"), sep=",")
with mock.patch("snowflake.connector.connect") as mock_connect:
with Snowflake(profile=PROFILE, dbapi=DBAPIS, **sf_credentials) as sf:
sf.to_dataframe()
sf.conn.cursor.return_value.fetch_pandas_all.assert_called_with()
sf.to_dataframe(5)
sf.conn.cursor.return_value.fetchmany.assert_called_with(5)
def test_with_connect(mock_session, sf_credentials):
with mock.patch("snowflake.connector.connect") as mock_connect:
with Snowflake(profile=PROFILE, dbapi=DBAPIS, **sf_credentials) as sf:
mock_connect.assert_called_with(
account="account",
user="******",
warehouse="warehouse",
password="******",
database="database",
schema="schema",
)
sf.conn.cursor.assert_called_with()
sf.conn.cursor.return_value.execute.assert_any_call("USE WAREHOUSE warehouse", ())
sf.conn.cursor.return_value.execute.assert_any_call("USE DATABASE database", ())
sf.conn.cursor.return_value.execute.assert_any_call("USE SCHEMA schema", ())
mock_connect.side_effect = Exception("Connect Exception")
with pytest.raises(DBError):
with Snowflake(profile=PROFILE, dbapi=DBAPIS, **sf_credentials) as sf:
sf.cursor
def test_constructor_yaml(input_kms_key, profile, mock_session):
sf = Snowflake(
profile=profile, kms_key=input_kms_key, dbapi=DBAPIS, config_yaml="some_config.yml"
)
mock_session.assert_called_with(profile_name=profile)
assert sf.profile == profile
assert sf.kms_key == input_kms_key
assert sf.connection["account"] == "account"
assert sf.connection["warehouse"] == "warehouse"
assert sf.connection["database"] == "database"
assert sf.connection["user"] == "user"
assert sf.connection["password"] == "password"
def test_constructor_yaml(mock_session):
sf = Snowflake(profile=PROFILE,
kms_key=KMS,
dbapi=DBAPIS,
config_yaml="some_config.yml")
mock_session.assert_called_with(profile_name=PROFILE)
assert sf.profile == PROFILE
assert sf.kms_key == KMS
assert sf.connection["account"] == "account"
assert sf.connection["warehouse"] == "warehouse"
assert sf.connection["database"] == "database"
assert sf.connection["user"] == "user"
assert sf.connection["password"] == "password"
def test_constructor(mock_session, sf_credentials):
sf = Snowflake(profile=PROFILE,
kms_key=KMS,
dbapi=DBAPIS,
**sf_credentials)
mock_session.assert_called_with(profile_name=PROFILE)
assert sf.profile == PROFILE
assert sf.kms_key == KMS
assert sf.connection["account"] == "account"
assert sf.connection["warehouse"] == "warehouse"
assert sf.connection["database"] == "database"
assert sf.connection["user"] == "user"
assert sf.connection["password"] == "password"
def test_unload_exception(mock_session, sf_credentials):
with mock.patch("snowflake.connector.connect") as mock_connect:
with Snowflake(profile=PROFILE, dbapi=DBAPIS, **sf_credentials) as sf:
with pytest.raises(ValueError):
sf.unload("table_name", "@~/stage", file_type="unknown")
# exception
sf.conn.cursor.return_value.execute.side_effect = Exception("UNLOAD Exception")
with pytest.raises(DBError):
sf.unload("@~/stage", "table_name")
sf.conn = None
with pytest.raises(DBError):
sf.unload("@~/stage", "table_name")
def test_download_from_internal(mock_session, sf_credentials):
with mock.patch("snowflake.connector.connect") as mock_connect:
with Snowflake(profile=PROFILE, dbapi=DBAPIS, **sf_credentials) as sf:
sf.download_from_internal("@~/internal", "/some/file")
sf.conn.cursor.return_value.execute.assert_called_with(
"GET @~/internal 'file:///some/file' PARALLEL=10", ())
sf.download_from_internal("@~/internal", "/some/file", parallel=99)
sf.conn.cursor.return_value.execute.assert_called_with(
"GET @~/internal 'file:///some/file' PARALLEL=99", ())
# exception
sf.conn.cursor.return_value.execute.side_effect = Exception(
"GET Exception")
with pytest.raises(DBError):
sf.download_from_internal("@~/internal", "/some/file")
