Repository: PrefectHQ/prefect

Issue:
Postgres Tasks
Using https://github.com/psycopg/psycopg2 let's add some tasks for working with a postgres database. The main function that we should add is `execute` which takes in an arbitrary SQL string to execute against a db.

On immediate appearance the `ExecuteTask` (or whatever it becomes called) should take in a db host, db name, user, pass, SQL string, data tuple, and optional commit param. http://initd.org/psycopg/docs/cursor.html?highlight=execute#cursor.execute

There are other functions such as `fetchone`, `fetchmany`, `rowcount`, etc... that may only be _slightly_ useful and we should use this issue as a means to debate what we think is important. (http://initd.org/psycopg/docs/cursor.html)

cc @jlowin for your psycopg2 expertise


Discussion:
@joshmeek did you want to add any other tasks besides execute?
Good question, I think `execute` is the most useful in it's ability to essentially allow any SQL command. If you see any others in `psycopg2` which seem crucial for a Postgres aspect of someone's workflow I welcome them but `execute` is a good first step!
I'll add in a separate task for `fetch` for now too and we can go from there

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -18,6 +18,7 @@
     ],
     "kubernetes": ["kubernetes >= 8.0.1, < 9.0", "dask-kubernetes == 0.7.0"],
     "rss": ["feedparser >= 5.0.1, < 6.0"],
+    "postgres": ["psycopg2-binary >= 2.8.2"],
     "spacy": ["spacy >= 2.0.0, < 3.0.0"],
     "templates": ["jinja2 >= 2.0, < 3.0"],
     "viz": ["graphviz >= 0.8.3"],
diff --git a/src/prefect/tasks/postgres/__init__.py b/src/prefect/tasks/postgres/__init__.py
new file mode 100644
--- /dev/null
+++ b/src/prefect/tasks/postgres/__init__.py
@@ -0,0 +1,11 @@
+"""
+This module contains a collection of tasks for interacting with Postgres databases via
+the psycopg2 library.
+"""
+
+try:
+    from prefect.tasks.postgres.postgres import PostgresExecute, PostgresFetch
+except ImportError:
+    raise ImportError(
+        'Using `prefect.tasks.postgres` requires Prefect to be installed with the "postgres" extra.'
+    )
diff --git a/src/prefect/tasks/postgres/postgres.py b/src/prefect/tasks/postgres/postgres.py
new file mode 100644
--- /dev/null
+++ b/src/prefect/tasks/postgres/postgres.py
@@ -0,0 +1,208 @@
+from prefect import Task
+from prefect.utilities.tasks import defaults_from_attrs
+
+import psycopg2 as pg
+
+
+class PostgresExecute(Task):
+    """
+        Task for executing a query against a Postgres database.
+
+        Args:
+            - db_name (str): name of Postgres database
+            - user (str): user name used to authenticate
+            - password (str): password used to authenticate
+            - host (str): database host address
+            - port (int, optional): port used to connect to Postgres database, defaults to 5432 if not provided
+            - query (str, optional): query to execute against database
+            - data (tuple, optional): values to use in query, must be specified using placeholder is query string
+            - commit (bool, optional): set to True to commit transaction, defaults to false
+            - **kwargs (dict, optional): additional keyword arguments to pass to the
+                Task constructor
+        """
+
+    def __init__(
+        self,
+        db_name: str,
+        user: str,
+        password: str,
+        host: str,
+        port: int = 5432,
+        query: str = None,
+        data: tuple = None,
+        commit: bool = False,
+        **kwargs
+    ):
+        self.db_name = db_name
+        self.user = user
+        self.password = password
+        self.host = host
+        self.port = port
+        self.query = query
+        self.data = data
+        self.commit = commit
+        super().__init__(**kwargs)
+
+    @defaults_from_attrs("query", "data", "commit")
+    def run(self, query: str = None, data: tuple = None, commit: bool = False):
+        """
+        Task run method. Executes a query against Postgres database.
+
+        Args:
+            - query (str, optional): query to execute against database
+            - data (tuple, optional): values to use in query, must be specified using
+                placeholder is query string
+            - commit (bool, optional): set to True to commit transaction, defaults to false
+
+        Returns:
+            - None
+
+        Raises:
+            - ValueError: if query parameter is None or a blank string
+            - DatabaseError: if exception occurs when executing the query
+        """
+        if not query:
+            raise ValueError("A query string must be provided")
+
+        ## connect to database, open cursor
+        ## allow psycopg2 to pass through any exceptions raised
+        conn = pg.connect(
+            dbname=self.db_name,
+            user=self.user,
+            password=self.password,
+            host=self.host,
+            port=self.port,
+        )
+
+        ## try to execute query
+        ## context manager automatically rolls back failed transactions
+        try:
+            with conn:
+                with conn.cursor() as cursor:
+                    executed = cursor.execute(query=query, vars=data)
+                    if commit:
+                        conn.commit()
+
+            conn.close()
+            return executed
+
+        ## pass through error, and ensure connection is closed
+        except (Exception, pg.DatabaseError) as error:
+            conn.close()
+            raise error
+
+
+class PostgresFetch(Task):
+    """
+        Task for fetching results of query from Postgres database.
+
+        Args:
+            - db_name (str): name of Postgres database
+            - user (str): user name used to authenticate
+            - password (str): password used to authenticate
+            - host (str): database host address
+            - port (int, optional): port used to connect to Postgres database, defaults to 5432 if not provided
+            - fetch (str, optional): one | many | all, used to determine how many results to fetch from executed query
+            - fetch_count (int, optional): if fetch = 'many', determines the number of results to fetch, defaults to 10
+            - query (str, optional): query to execute against database
+            - data (tuple, optional): values to use in query, must be specified using placeholder is query string
+            - commit (bool, optional): set to True to commit transaction, defaults to false
+            - **kwargs (dict, optional): additional keyword arguments to pass to the
+                Task constructor
+        """
+
+    def __init__(
+        self,
+        db_name: str,
+        user: str,
+        password: str,
+        host: str,
+        port: int = 5432,
+        fetch: str = "one",
+        fetch_count: int = 10,
+        query: str = None,
+        data: tuple = None,
+        commit: bool = False,
+        **kwargs
+    ):
+        self.db_name = db_name
+        self.user = user
+        self.password = password
+        self.host = host
+        self.port = port
+        self.fetch = fetch
+        self.fetch_count = fetch_count
+        self.query = query
+        self.data = data
+        self.commit = commit
+        super().__init__(**kwargs)
+
+    @defaults_from_attrs("fetch", "fetch_count", "query", "data", "commit")
+    def run(
+        self,
+        fetch: str = "one",
+        fetch_count: int = 10,
+        query: str = None,
+        data: tuple = None,
+        commit: bool = False,
+    ):
+        """
+        Task run method. Executes a query against Postgres database and fetches results.
+
+        Args:
+            - fetch (str, optional): one | many | all, used to determine how many results to fetch from executed query
+            - fetch_count (int, optional): if fetch = 'many', determines the number of results to fetch, defaults to 10
+            - query (str, optional): query to execute against database
+            - data (tuple, optional): values to use in query, must be specified using placeholder is query string
+            - commit (bool, optional): set to True to commit transaction, defaults to false
+
+        Returns:
+            - records (tuple or list of tuples): records from provided query
+
+        Raises:
+            - ValueError: if query parameter is None or a blank string
+            - DatabaseError: if exception occurs when executing the query
+        """
+        if not query:
+            raise ValueError("A query string must be provided")
+
+        if fetch not in {"one", "many", "all"}:
+            raise ValueError(
+                "The 'fetch' parameter must be one of the following - ('one', 'many', 'all')"
+            )
+
+        ## connect to database, open cursor
+        ## allow psycopg2 to pass through any exceptions raised
+        conn = pg.connect(
+            dbname=self.db_name,
+            user=self.user,
+            password=self.password,
+            host=self.host,
+            port=self.port,
+        )
+
+        ## try to execute query
+        ## context manager automatically rolls back failed transactions
+        try:
+            with conn:
+                with conn.cursor() as cursor:
+                    cursor.execute(query=query, vars=data)
+
+                    ## fetch results
+                    if fetch == "all":
+                        records = cursor.fetchall()
+                    elif fetch == "many":
+                        records = cursor.fetchmany(fetch_count)
+                    else:
+                        records = cursor.fetchone()
+
+                    if commit:
+                        conn.commit()
+
+            conn.close()
+            return records
+
+        ## pass through error, and ensure connection is closed
+        except (Exception, pg.DatabaseError) as error:
+            conn.close()
+            raise error


Test patch:
diff --git a/tests/tasks/postgres/__init__.py b/tests/tasks/postgres/__init__.py
new file mode 100644
--- /dev/null
+++ b/tests/tasks/postgres/__init__.py
@@ -0,0 +1,3 @@
+import pytest
+
+pytest.importorskip("psycopg2")
diff --git a/tests/tasks/postgres/test_postgres.py b/tests/tasks/postgres/test_postgres.py
new file mode 100644
--- /dev/null
+++ b/tests/tasks/postgres/test_postgres.py
@@ -0,0 +1,40 @@
+import pytest
+
+from prefect.tasks.postgres import PostgresExecute, PostgresFetch
+
+
+class TestPostgresExecute:
+    def test_construction(self):
+        task = PostgresExecute(
+            db_name="test", user="test", password="test", host="test"
+        )
+        assert task.commit is False
+
+    def test_query_string_must_be_provided(self):
+        task = PostgresExecute(
+            db_name="test", user="test", password="test", host="test"
+        )
+        with pytest.raises(ValueError) as exc:
+            task.run()
+        assert "A query string must be provided" == str(exc.value)
+
+
+class TestPostgresFetch:
+    def test_construction(self):
+        task = PostgresFetch(db_name="test", user="test", password="test", host="test")
+        assert task.fetch == "one"
+
+    def test_query_string_must_be_provided(self):
+        task = PostgresFetch(db_name="test", user="test", password="test", host="test")
+        with pytest.raises(ValueError) as exc:
+            task.run()
+        assert "A query string must be provided" == str(exc.value)
+
+    def test_bad_fetch_param_raises(self):
+        task = PostgresFetch(db_name="test", user="test", password="test", host="test")
+        with pytest.raises(ValueError) as exc:
+            task.run(query="SELECT * FROM some_table", fetch="not a valid parameter")
+        assert (
+            "The 'fetch' parameter must be one of the following - ('one', 'many', 'all')"
+            == str(exc.value)
+        )