def test_execute_error_must_pass_through(self, monkeypatch):
snowflake_module_connect_method = MagicMock()
connection = MagicMock(spec=sf.SnowflakeConnection)
cursor = MagicMock(spec=sf.DictCursor)
# link all the mocks together appropriately
snowflake_module_connect_method.return_value = connection
connection.cursor = cursor
# database cursors can be ugly to mock given the use of __enter__
cursor.return_value.__enter__.return_value.execute.side_effect = (
sf.DatabaseError("Invalid query"))
snowflake_connector_module = MagicMock(
connect=snowflake_module_connect_method)
monkeypatch.setattr("prefect.tasks.snowflake.snowflake.sf",
snowflake_connector_module)
task = SnowflakeQuery(account="test",
user="******",
password="******",
warehouse="test")
with pytest.raises(sf.errors.DatabaseError, match="Invalid query"):
task.run(query="SELECT * FROM foo")
def test_query_string_must_be_provided(self):
task = SnowflakeQuery(account="test",
user="******",
password="******",
warehouse="test")
with pytest.raises(ValueError,
match="A query string must be provided"):
task.run()
def test_required_parameters(self):
# missing account
with pytest.raises(ValueError):
SnowflakeQuery().run(user="******", password="******")
# missing user
with pytest.raises(ValueError):
SnowflakeQuery().run(account="test", password="******")
# missing query
with pytest.raises(ValueError):
SnowflakeQuery().run(account="test", user="******", password="******")
def test_run_method_accepts_alternate_cursor(self, monkeypatch):
"""
Tests that the default cursor used by the Snowflake connection
is of the type `SnowflakeCursor`, and that the cursor type can
be overriden by providing a different cursor type to the
`cursor_type` keyword argument.
"""
snowflake_module_connect_method = MagicMock()
connection = MagicMock(spec=sf.SnowflakeConnection)
default_cursor = MagicMock(spec=sf.cursor.SnowflakeCursor)
dict_cursor = MagicMock(spec=sf.cursor.DictCursor)
# link all the mocks together appropriately
snowflake_module_connect_method.return_value = connection
# setting fetchall return
default_cursor.__enter__.return_value.execute.return_value.fetchall.return_value = [
"1",
"2",
"3",
]
dict_cursor.__enter__.return_value.execute.return_value.fetchall.return_value = {
"one": "1",
"two": "2",
"three": "3",
}
snowflake_connector_module = MagicMock(
connect=snowflake_module_connect_method)
def mock_cursor(cursor_class: SnowflakeCursor) -> SnowflakeCursor:
if cursor_class == DictCursor:
return dict_cursor
else:
return default_cursor
connection.cursor = mock_cursor
monkeypatch.setattr("prefect.tasks.snowflake.snowflake.sf",
snowflake_connector_module)
query = "select * from numbers"
task = SnowflakeQuery(account="test",
user="******",
password="******",
query=query)
default_cursor_output = task.run(cursor_type=SnowflakeCursor)
dict_cursor_output = task.run(cursor_type=DictCursor)
assert default_cursor_output == ["1", "2", "3"]
assert dict_cursor_output == {"one": "1", "two": "2", "three": "3"}
def test_runtime(self, monkeypatch):
cursor = MagicMock(spec=sf.DictCursor)
connection = MagicMock(spec=sf.SnowflakeConnection, cursor=cursor)
snowflake_module_connect_method = MagicMock(return_value=connection)
snowflake_connector_module = MagicMock(connect=snowflake_module_connect_method)
# setting fetchall return
cursor.return_value.__enter__.return_value.execute.return_value.fetchall.return_value = [
"TESTDB"
]
monkeypatch.setattr(
"prefect.tasks.snowflake.snowflake.sf", snowflake_connector_module
)
# task needs to allow for purely runtime arguments
output = SnowflakeQuery().run(
account="test", user="******", password="******", query="SHOW DATABASES"
)
assert output == ["TESTDB"]
for call in snowflake_module_connect_method.call_args_list:
args, kwargs = call
assert kwargs == dict(
account="test",
user="******",
password="******",
application="Prefect_SnowflakeQuery",
)
def test_execute_fetchall(self, monkeypatch):
"""
Tests that the SnowflakeQuery Task calls the fetchall method on the
cursor. This is to prevent future code edits from returning the cursor
object because that cursors are not pickleable.
"""
snowflake_module_connect_method = MagicMock()
connection = MagicMock(spec=sf.SnowflakeConnection)
cursor = MagicMock(spec=sf.DictCursor)
# link all the mocks together appropriately
snowflake_module_connect_method.return_value = connection
connection.cursor = cursor
# setting fetchall return
cursor.return_value.__enter__.return_value.execute.return_value.fetchall.return_value = [
"TESTDB"
]
snowflake_connector_module = MagicMock(connect=snowflake_module_connect_method)
monkeypatch.setattr(
"prefect.tasks.snowflake.snowflake.sf", snowflake_connector_module
)
query = "SHOW DATABASES"
output = SnowflakeQuery(
account="test", user="******", password="******", query=query
).run()
assert output == ["TESTDB"]
def test_construction(self):
task = SnowflakeQuery(account="test",
user="******",
password="******",
warehouse="test")
assert task.autocommit is None
assert task.authenticator is None
