def test_query_string_must_be_provided(self):
task = MySQLFetch(db_name="test",
user="******",
password="******",
host="test")
with pytest.raises(ValueError,
match="A query string must be provided"):
task.run()
def test_run_args_used_over_init_args(self, monkeypatch):
mock_connect = MagicMock()
monkeypatch.setattr("pymysql.connect", mock_connect)
task = MySQLFetch(db_name="test",
user="******",
password="******",
host="test")
task.run(query="select * from users", password="******")
assert mock_connect.call_args[1]["password"] == "password_from_secret"
def test_bad_cursor_type_param_type_raises(self):
task = MySQLFetch(db_name="test",
user="******",
password="******",
host="test")
with pytest.raises(
TypeError,
match=
f"'cursor_type' should be one of \['cursor', 'dictcursor', 'sscursor', 'ssdictcursor'\] or a full cursor class, got \['cursor'\]",
):
task.run(query="SELECT * FROM some_table", cursor_type=["cursor"])
def test_bad_fetch_param_raises(self):
task = MySQLFetch(db_name="test",
user="******",
password="******",
host="test")
with pytest.raises(
ValueError,
match=
"The 'fetch' parameter must be one of the following - \('one', 'many', 'all'\)",
):
task.run(query="SELECT * FROM some_table",
fetch="not a valid parameter")
def test_unsupported_cursor_type_str_param_raises(self):
cursor_type = "unsupportedcursor"
task = MySQLFetch(db_name="test",
user="******",
password="******",
host="test")
with pytest.raises(
TypeError,
match=
rf"'cursor_type' should be one of \['cursor', 'dictcursor', 'sscursor', 'ssdictcursor'\] or a full cursor class, got {cursor_type}",
):
task.run(query="SELECT * FROM some_table", cursor_type=cursor_type)
def test_construction_with_cursor_type_class(self):
task = MySQLFetch(
db_name="test",
user="******",
password="******",
host="test",
cursor_type=pymysql.cursors.DictCursor,
)
assert task.cursor_type == pymysql.cursors.DictCursor
def test_construction_with_cursor_type_str(self):
task = MySQLFetch(
db_name="test",
user="******",
password="******",
host="test",
cursor_type="dictcursor",
)
assert task.cursor_type == "dictcursor"
def test_construction(self):
task = MySQLFetch(db_name="test",
user="******",
password="******",
host="test")
assert task.fetch == "one"
from prefect.tasks.mysql.mysql import MySQLFetch, MySQLExecute
from prefect import Flow, task
EXAMPLE_TABLE = "user"
HOST = "localhost"
PORT = 3306
DB_NAME = "ext"
USER = "******"
PASSWORD = "******"
mysql_fetch = MySQLFetch(host=HOST,
port=PORT,
db_name=DB_NAME,
user=USER,
password=PASSWORD)
mysql_exec = MySQLExecute(host=HOST,
port=PORT,
db_name=DB_NAME,
user=USER,
password=PASSWORD)
@task
def print_results(x):
print(x)
with Flow("MySQL Example") as flow:
# fetch 3 results
fetch_results = mysql_fetch(query=f"SELECT * FROM {EXAMPLE_TABLE}",
