Repository: PrefectHQ/prefect

Issue:
Discussion: `SqliteTask`
I am currently working on a simple implementation of a `SqliteTask` for interacting with `sqlite3` databases.

What I'd like to support:
- executing single statements and returning the values of `fetchall()` (or some prespecified number of return rows)
- executing full sql scripts (either as a string or read from a file); note that if you execute a script only, there is nothing to return

My current sketch design idea:
```python
class SqliteTask(Task):
    def __init__(self, statement: str = None, 
                       script: str = None, db: str = ":memory:", 
                       return_rows: int = 0, **kwargs):
        ...
```
Moreover, if you provide _both_ `script` and `statement`, the `script` will be executed first and the `statement` last, in case it's a `SELECT` statement to query what just happened in the script.

I'm opening this issue for discussion in case this is overloading the task with too many parameters / options / possibilities.  I want to make sure it's still clear as a single `Task`.

**cc**: @jlowin @joshmeek 


Discussion:

Resolving patch:
diff --git a/src/prefect/tasks/__init__.py b/src/prefect/tasks/__init__.py
--- a/src/prefect/tasks/__init__.py
+++ b/src/prefect/tasks/__init__.py
@@ -3,5 +3,6 @@
 from prefect.core.task import Task
 import prefect.tasks.core
 import prefect.tasks.control_flow
+import prefect.tasks.database
 import prefect.tasks.shell
 import prefect.tasks.templates
diff --git a/src/prefect/tasks/database/__init__.py b/src/prefect/tasks/database/__init__.py
new file mode 100644
--- /dev/null
+++ b/src/prefect/tasks/database/__init__.py
@@ -0,0 +1,3 @@
+# Licensed under LICENSE.md; also available at https://www.prefect.io/licenses/alpha-eula
+
+from prefect.tasks.database.sqlite import SQLiteQueryTask, SQLiteScriptTask
diff --git a/src/prefect/tasks/database/sqlite.py b/src/prefect/tasks/database/sqlite.py
new file mode 100644
--- /dev/null
+++ b/src/prefect/tasks/database/sqlite.py
@@ -0,0 +1,74 @@
+# Licensed under LICENSE.md; also available at https://www.prefect.io/licenses/alpha-eula
+
+import sqlite3 as sql
+from contextlib import closing
+
+import prefect
+from prefect import Task
+
+
+class SQLiteQueryTask(Task):
+    """
+    Task for executing a single query against a sqlite3 database; returns
+    the result (if any) from the query.
+
+    Args:
+        - db (str): the location of the database (.db) file
+        - query (str, optional): the optional _default_ query to execute at runtime;
+            can also be provided as a keyword to `run`, which takes precendence over this default.
+            Note that a query should consist of a _single SQL statement_.
+        - **kwargs (optional): additional keyword arguments to pass to the
+            standard Task initalization
+    """
+
+    def __init__(self, db: str, query: str = None, **kwargs) -> None:
+        self.db = db
+        self.query = query
+        super().__init__(**kwargs)
+
+    def run(self, query: str = None):
+        """
+        Args:
+            - query (str, optional): the optional query to execute at runtime;
+                if not provided, `self.query` will be used instead. Note that a query should consist of a _single SQL statement_.
+
+        Returns:
+            - [Any]: the results of the query
+        """
+        query = self.query if query is None else query
+        with closing(sql.connect(self.db)) as conn:
+            with closing(conn.cursor()) as cursor:
+                cursor.execute(query)
+                out = cursor.fetchall()
+                conn.commit()
+        return out
+
+
+class SQLiteScriptTask(Task):
+    """
+    Task for executing a SQL script against a sqlite3 database.
+
+    Args:
+        - db (str): the location of the database (.db) file
+        - script (str, optional): the optional _default_ script string to render at runtime;
+            can also be provided as a keyword to `run`, which takes precendence over this default.
+        - **kwargs (optional): additional keyword arguments to pass to the
+            standard Task initialization
+    """
+
+    def __init__(self, db: str, script: str = None, **kwargs) -> None:
+        self.db = db
+        self.script = script
+        super().__init__(**kwargs)
+
+    def run(self, script: str = None):
+        """
+        Args:
+            - script (str, optional): the optional script to execute at runtime;
+                if not provided, `self.script` will be used instead.
+        """
+        script = self.script if script is None else script
+        with closing(sql.connect(self.db)) as conn:
+            with closing(conn.cursor()) as cursor:
+                cursor.executescript(script)
+                conn.commit()


Test patch:
diff --git a/tests/tasks/test_sqlite.py b/tests/tasks/test_sqlite.py
new file mode 100644
--- /dev/null
+++ b/tests/tasks/test_sqlite.py
@@ -0,0 +1,101 @@
+import pytest
+import sqlite3
+import tempfile
+from contextlib import closing
+
+from prefect import Flow
+from prefect.tasks.database import SQLiteQueryTask, SQLiteScriptTask
+from prefect.utilities.tests import raise_on_exception
+
+
+sql_script = """
+CREATE TABLE TEST (NUMBER INTEGER, DATA TEXT);
+INSERT INTO TEST (NUMBER, DATA) VALUES
+(11, 'first'),
+(12, 'second'),
+(13, 'third');
+"""
+
+
+@pytest.fixture(scope="module")
+def database():
+    with tempfile.NamedTemporaryFile(suffix=".db") as tmp:
+        with closing(sqlite3.connect(tmp.name)) as conn:
+            with closing(conn.cursor()) as c:
+                c.executescript(sql_script)
+                conn.commit()
+        yield tmp.name
+
+
+class TestSQLiteQueryTask:
+    def test_sqlite_query_task_requires_db(self):
+        with pytest.raises(TypeError):
+            task = SQLiteQueryTask()
+
+    def test_sqlite_query_task_initializes_and_runs_basic_query(self, database):
+        with Flow() as f:
+            task = SQLiteQueryTask(db=database)(query="SELECT * FROM TEST")
+        out = f.run(return_tasks=[task])
+        assert out.is_successful()
+        assert out.result[task].result == [(11, "first"), (12, "second"), (13, "third")]
+
+    def test_sqlite_query_task_initializes_with_query_and_runs(self, database):
+        with Flow() as f:
+            task = SQLiteQueryTask(db=database, query="SELECT * FROM TEST")()
+        out = f.run(return_tasks=[task])
+        assert out.is_successful()
+        assert out.result[task].result == [(11, "first"), (12, "second"), (13, "third")]
+
+    def test_sqlite_error_results_in_failed_state(self, database):
+        with Flow() as f:
+            task = SQLiteQueryTask(db=database, query="SELECT * FROM FOOBAR")()
+        out = f.run(return_tasks=[task])
+        assert out.is_failed()
+        assert "no such table: FOOBAR" in str(out.result[task].result)
+
+    def test_only_single_statement_queries_allowed(self, database):
+        query = """INSERT INTO TEST (NUMBER, DATA) VALUES\n(88, "other");\nSELECT * FROM TEST;"""
+        with Flow() as f:
+            task = SQLiteQueryTask(db=database, query=query)()
+        out = f.run(return_tasks=[task])
+        assert out.is_failed()
+        assert "one statement at a time" in str(out.result[task].result)
+
+
+class TestSQLiteScriptTask:
+    def test_sqlite_script_task_requires_db(self):
+        with pytest.raises(TypeError):
+            task = SQLiteScriptTask()
+
+    def test_sqlite_script_task_initializes_and_runs_basic_script(self, database):
+        with Flow() as f:
+            task = SQLiteScriptTask(db=database)(script="SELECT * FROM TEST;")
+        out = f.run(return_tasks=[task])
+        assert out.is_successful()
+        assert out.result[task].result is None
+
+    def test_sqlite_script_task_initializes_with_script_and_runs(self, database):
+        with Flow() as f:
+            task = SQLiteScriptTask(db=database, script="SELECT * FROM TEST;")()
+        out = f.run(return_tasks=[task])
+        assert out.is_successful()
+        assert out.result[task].result is None
+
+    def test_sqlite_error_results_in_failed_state(self, database):
+        with Flow() as f:
+            task = SQLiteScriptTask(db=database)(script="SELECT * FROM FOOBAR;")
+        out = f.run(return_tasks=[task])
+        assert out.is_failed()
+        assert "no such table: FOOBAR" in str(out.result[task].result)
+
+
+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() as f:
+        alter = SQLiteScriptTask(db=database)(script=script)
+        task = SQLiteQueryTask(db=database, query="SELECT * FROM TEST2")(
+            upstream_tasks=[alter]
+        )
+    out = f.run(return_tasks=[task])
+    assert out.is_successful()
+    assert out.result[task].result == [(88, "other", None)]