def test_query_parameterization_data_on_init(self, database):
with Flow(name="test") as f:
task = SQLiteQuery(db=database, data=(12, 13))(
query="SELECT * FROM TEST WHERE NUMBER in (?, ?)")
out = f.run()
assert out.is_successful()
assert out.result[task].result == [(12, "second"), (13, "third")]
def test_only_single_statement_queries_allowed(self, database):
query = """INSERT INTO TEST (NUMBER, DATA) VALUES\n(88, "other");\nSELECT * FROM TEST;"""
with Flow(name="test") as f:
task = SQLiteQuery(db=database, query=query)()
out = f.run()
assert out.is_failed()
assert "one statement at a time" in str(out.result[task].result)
def test_sqlite_query_task_initializes_with_query_and_runs(self, database):
with Flow(name="test") as f:
task = SQLiteQuery(db=database, query="SELECT * FROM TEST")()
out = f.run()
assert out.is_successful()
assert out.result[task].result == [(11, "first"), (12, "second"),
(13, "third")]
def test_unparametrized_query_no_data(self, database):
with Flow(name="test") as f:
task = SQLiteQuery(db=database)(query="SELECT * FROM TEST;",
data=())
out = f.run()
assert out.is_successful()
assert out.result[task].result == [(11, "first"), (12, "second"),
(13, "third")]
def test_composition_of_tasks(database):
script = """CREATE TABLE TEST2 (NUM INTEGER, DATA TEXT); INSERT INTO TEST2 (NUM, DATA) VALUES\n(88, "other"); ALTER TABLE TEST2\n ADD status TEXT;"""
with Flow(name="test") as f:
alter = SQLiteScript(db=database)(script=script)
task = SQLiteQuery(db=database,
query="SELECT * FROM TEST2")(upstream_tasks=[alter])
out = f.run()
assert out.is_successful()
assert out.result[task].result == [(88, "other", None)]
def test_parametrization_of_tasks(database):
with Flow(name="test") as f:
db = Parameter("db")
script = Parameter("script")
script = SQLiteScript(db=database)(script=script)
task = SQLiteQuery()(db=db,
query="SELECT * FROM TEST WHERE NUMBER = 14",
upstream_tasks=[script])
out = f.run(db=database,
script="INSERT INTO TEST (NUMBER, DATA) VALUES (14, 'fourth')")
assert out.is_successful()
assert out.result[task].result == [(14, "fourth")]
def test_sqlite_error_results_in_failed_state(self, database):
with Flow(name="test") as f:
task = SQLiteQuery(db=database, query="SELECT * FROM FOOBAR")()
out = f.run()
assert out.is_failed()
assert "no such table: FOOBAR" in str(out.result[task].result)
def test_sqlite_query_task_requires_db(self):
with pytest.raises(TypeError):
task = SQLiteQuery()
def test_initialization(self):
task = SQLiteQuery()
from prefect import task, Flow, Parameter
from prefect.tasks.database import SQLiteQuery
from prefect.tasks.redis import RedisSet
@task
def a(x):
print(x)
with Flow("meta_model_flow") as meta_model_flow:
db = Parameter("db")
# a(db)
models = SQLiteQuery(db=db, query="SELECT identifier FROM models")()
meta_model_flow.add_task(models)
models.set_upstream(db, key=db)
meta_model_flow.run(db="test")