[AIRFLOW-4734] Upsert functionality for PostgresHook.insert_rows() (#8625)

PostgresHook's parent class, DbApiHook, implements upsert in its insert_rows() method
with the replace=True flag. However, the underlying generated SQL is specific to MySQL's
"REPLACE INTO" syntax and is not applicable to PostgreSQL.

This pulls out the sql generation code for insert/upsert out in to a method that is then
overridden in the PostgreSQL subclass to generate the "INSERT ... ON CONFLICT DO
UPDATE" syntax ("new" since Postgres 9.5)

Author: oxymor0n

airflow/hooks/dbapi_hook.py

@@ -226,8 +226,43 @@ def get_cursor(self):
         """
         return self.get_conn().cursor()
 
+    @staticmethod
+    def _generate_insert_sql(table, values, target_fields, replace, **kwargs):
+        """
+        Static helper method that generate the INSERT SQL statement.
+        The REPLACE variant is specific to MySQL syntax.
+
+        :param table: Name of the target table
+        :type table: str
+        :param values: The row to insert into the table
+        :type values: tuple of cell values
+        :param target_fields: The names of the columns to fill in the table
+        :type target_fields: iterable of strings
+        :param replace: Whether to replace instead of insert
+        :type replace: bool
+        :return: The generated INSERT or REPLACE SQL statement
+        :rtype: str
+        """
+        placeholders = ["%s", ] * len(values)
+
+        if target_fields:
+            target_fields = ", ".join(target_fields)
+            target_fields = "({})".format(target_fields)
+        else:
+            target_fields = ''
+
+        if not replace:
+            sql = "INSERT INTO "
+        else:
+            sql = "REPLACE INTO "
+        sql += "{0} {1} VALUES ({2})".format(
+            table,
+            target_fields,
+            ",".join(placeholders))
+        return sql
+
     def insert_rows(self, table, rows, target_fields=None, commit_every=1000,
-                    replace=False):
+                    replace=False, **kwargs):
         """
         A generic way to insert a set of tuples into a table,
         a new transaction is created every commit_every rows
@@ -244,11 +279,6 @@ def insert_rows(self, table, rows, target_fields=None, commit_every=1000,
         :param replace: Whether to replace instead of insert
         :type replace: bool
         """
-        if target_fields:
-            target_fields = ", ".join(target_fields)
-            target_fields = "({})".format(target_fields)
-        else:
-            target_fields = ''
         i = 0
         with closing(self.get_conn()) as conn:
             if self.supports_autocommit:
@@ -262,15 +292,9 @@ def insert_rows(self, table, rows, target_fields=None, commit_every=1000,
                     for cell in row:
                         lst.append(self._serialize_cell(cell, conn))
                     values = tuple(lst)
-                    placeholders = ["%s", ] * len(values)
-                    if not replace:
-                        sql = "INSERT INTO "
-                    else:
-                        sql = "REPLACE INTO "
-                    sql += "{0} {1} VALUES ({2})".format(
-                        table,
-                        target_fields,
-                        ",".join(placeholders))
+                    sql = self._generate_insert_sql(
+                        table, values, target_fields, replace, **kwargs
+                    )
                     cur.execute(sql, values)
                     if commit_every and i % commit_every == 0:
                         conn.commit()

airflow/providers/google/cloud/hooks/bigquery.py

@@ -119,7 +119,8 @@ def get_client(self, project_id: Optional[str] = None, location: Optional[str] =
         )
 
     def insert_rows(
-        self, table: Any, rows: Any, target_fields: Any = None, commit_every: Any = 1000, replace: Any = False
+        self, table: Any, rows: Any, target_fields: Any = None, commit_every: Any = 1000,
+        replace: Any = False, **kwargs
     ) -> NoReturn:
         """
         Insertion is currently unsupported. Theoretically, you could use

airflow/providers/postgres/hooks/postgres.py

@@ -180,3 +180,57 @@ def get_iam_token(self, conn):
         else:
             token = aws_hook.conn.generate_db_auth_token(conn.host, port, conn.login)
         return login, token, port
+
+    @staticmethod
+    def _generate_insert_sql(table, values, target_fields, replace, **kwargs):
+        """
+        Static helper method that generate the INSERT SQL statement.
+        The REPLACE variant is specific to MySQL syntax.
+
+        :param table: Name of the target table
+        :type table: str
+        :param values: The row to insert into the table
+        :type values: tuple of cell values
+        :param target_fields: The names of the columns to fill in the table
+        :type target_fields: iterable of strings
+        :param replace: Whether to replace instead of insert
+        :type replace: bool
+        :param replace_index: the column or list of column names to act as
+            index for the ON CONFLICT clause
+        :type replace_index: str or list
+        :return: The generated INSERT or REPLACE SQL statement
+        :rtype: str
+        """
+        placeholders = ["%s", ] * len(values)
+        replace_index = kwargs.get("replace_index", None)
+
+        if target_fields:
+            target_fields_fragment = ", ".join(target_fields)
+            target_fields_fragment = "({})".format(target_fields_fragment)
+        else:
+            target_fields_fragment = ''
+
+        sql = "INSERT INTO {0} {1} VALUES ({2})".format(
+            table,
+            target_fields_fragment,
+            ",".join(placeholders))
+
+        if replace:
+            if target_fields is None:
+                raise ValueError("PostgreSQL ON CONFLICT upsert syntax requires column names")
+            if replace_index is None:
+                raise ValueError("PostgreSQL ON CONFLICT upsert syntax requires an unique index")
+            if isinstance(replace_index, str):
+                replace_index = [replace_index]
+            replace_index_set = set(replace_index)
+
+            replace_target = [
+                "{0} = excluded.{0}".format(col)
+                for col in target_fields
+                if col not in replace_index_set
+            ]
+            sql += " ON CONFLICT ({0}) DO UPDATE SET {1}".format(
+                ", ".join(replace_index),
+                ", ".join(replace_target),
+            )
+        return sql

tests/providers/postgres/hooks/test_postgres.py

@@ -199,3 +199,62 @@ def test_bulk_dump(self):
                     results = [line.rstrip().decode("utf-8") for line in f.readlines()]
 
         self.assertEqual(sorted(input_data), sorted(results))
+
+    @pytest.mark.backend("postgres")
+    def test_insert_rows(self):
+        table = "table"
+        rows = [("hello",),
+                ("world",)]
+
+        self.db_hook.insert_rows(table, rows)
+
+        assert self.conn.close.call_count == 1
+        assert self.cur.close.call_count == 1
+
+        commit_count = 2  # The first and last commit
+        self.assertEqual(commit_count, self.conn.commit.call_count)
+
+        sql = "INSERT INTO {}  VALUES (%s)".format(table)
+        for row in rows:
+            self.cur.execute.assert_any_call(sql, row)
+
+    @pytest.mark.backend("postgres")
+    def test_insert_rows_replace(self):
+        table = "table"
+        rows = [(1, "hello",),
+                (2, "world",)]
+        fields = ("id", "value")
+
+        self.db_hook.insert_rows(
+            table, rows, fields, replace=True, replace_index=fields[0])
+
+        assert self.conn.close.call_count == 1
+        assert self.cur.close.call_count == 1
+
+        commit_count = 2  # The first and last commit
+        self.assertEqual(commit_count, self.conn.commit.call_count)
+
+        sql = "INSERT INTO {0} ({1}, {2}) VALUES (%s,%s) " \
+              "ON CONFLICT ({1}) DO UPDATE SET {2} = excluded.{2}".format(
+                  table, fields[0], fields[1])
+        for row in rows:
+            self.cur.execute.assert_any_call(sql, row)
+
+    @pytest.mark.xfail
+    @pytest.mark.backend("postgres")
+    def test_insert_rows_replace_missing_target_field_arg(self):
+        table = "table"
+        rows = [(1, "hello",),
+                (2, "world",)]
+        fields = ("id", "value")
+        self.db_hook.insert_rows(
+            table, rows, replace=True, replace_index=fields[0])
+
+    @pytest.mark.xfail
+    @pytest.mark.backend("postgres")
+    def test_insert_rows_replace_missing_replace_index_arg(self):
+        table = "table"
+        rows = [(1, "hello",),
+                (2, "world",)]
+        fields = ("id", "value")
+        self.db_hook.insert_rows(table, rows, fields, replace=True)