BFD-4686: IDR Pipeline incremental performance improvements (#3125)

Author: malessi

apps/bfd-pipeline-idr/loader.py

@@ -4,6 +4,7 @@
 
 import psycopg
 from psycopg.abc import Params, QueryNoTemplate
+from psycopg.errors import DeadlockDetected, LockNotAvailable, QueryCanceled
 
 from constants import DEFAULT_MIN_DATE
 from load_partition import LoadPartition, LoadType
@@ -95,10 +96,14 @@ def __init__(
         self.meta_keys = (
             ["bfd_created_ts"] if self.immutable else ["bfd_created_ts", "bfd_updated_ts"]
         )
+        self.progress_start_timer = Timer("progress_start", model, partition)
         self.idr_query_timer = Timer("idr_query", model, partition)
         self.temp_table_timer = Timer("temp_table", model, partition)
         self.copy_timer = Timer("copy", model, partition)
-        self.insert_timer = Timer("insert", model, partition)
+        self.upsert_timer = Timer("upsert", model, partition)
+        self.last_updated_timer = Timer("last_updated", model, partition)
+        self.total_insert_timer = Timer("total_insert", model, partition)
+        self.update_progress_timer = Timer("update_progress", model, partition)
         self.commit_timer = Timer("commit", model, partition)
         self.load_type = load_type
         self.enable_load_progress = should_track_load_progress(load_mode)
@@ -111,8 +116,10 @@ def load(
         # (temp tables can't be created with an explicit schema set)
 
         with self.conn.cursor() as cur:
+            self.progress_start_timer.start()
             self._insert_batch_start(cur)
             self.conn.commit()
+            self.progress_start_timer.stop()
             data_loaded = False
             num_rows = 0
 
@@ -140,11 +147,13 @@ def load(
 
                 if results:
                     # Upsert into the main table
-                    self.insert_timer.start()
+                    self.total_insert_timer.start()
                     self._merge(cur, timestamp)
-                    self.insert_timer.stop()
+                    self.total_insert_timer.stop()
 
+                    self.update_progress_timer.start()
                     self._calculate_load_progress(cur, results)
+                    self.update_progress_timer.stop()
 
                 self.commit_timer.start()
                 self.conn.commit()
@@ -211,11 +220,13 @@ def _setup_temp_table(self, cur: psycopg.Cursor) -> None:
         # For simplicity's sake, we'll create our temp tables using the existing schema and
         # just drop the columns we need to ignore.
         cur.execute(
-            f"CREATE TEMPORARY TABLE {self.temp_table} (LIKE {self.table}) ON COMMIT DROP"  # type: ignore
+            f"CREATE TEMPORARY TABLE IF NOT EXISTS {self.temp_table} (LIKE {self.table}) "
+            "ON COMMIT PRESERVE ROWS"  # type: ignore
         )
+        cur.execute(f"TRUNCATE TABLE {self.temp_table}")  # type: ignore
         # Created/updated columns don't need to be loaded from the source.
         for col in self.meta_keys:
-            cur.execute(f"ALTER TABLE {self.temp_table} DROP COLUMN {col}")  # type: ignore
+            cur.execute(f"ALTER TABLE {self.temp_table} DROP COLUMN IF EXISTS {col}")  # type: ignore
 
     def _calculate_load_progress(self, cur: psycopg.Cursor, results: Sequence[T]) -> None:
         last = results[len(results) - 1].model_dump()
@@ -260,6 +271,7 @@ def _update_load_progress(
     ) -> None:
         if self.enable_load_progress:
             cur.execute(query, params)  # type: ignore
+            self.conn.commit()
 
     def _merge(self, cur: psycopg.Cursor, timestamp: datetime) -> None:
         unique_key = self.model.unique_key()
@@ -274,11 +286,15 @@ def _merge(self, cur: psycopg.Cursor, timestamp: datetime) -> None:
         on_conflict = (
             "DO NOTHING"
             if self.immutable or not update_set
-            else f"DO UPDATE SET {update_set}, bfd_updated_ts=%(timestamp)s"
+            else (
+                f"DO UPDATE SET {update_set}, bfd_updated_ts=%(timestamp)s "
+                "WHERE (t.*) IS DISTINCT FROM (EXCLUDED.*)"
+            )
         )
-        timestamp_placeholders = ",".join("%(timestamp)s" for _ in self.meta_keys)
+        timestamp_placeholders = ", ".join("%(timestamp)s" for _ in self.meta_keys)
 
         # Upsert into the main table
+        self.upsert_timer.start()
         if self.model.should_replace():
             # Delete before inserting since we've specified that the data should be
             # replaced rather than merged.
@@ -287,42 +303,54 @@ def _merge(self, cur: psycopg.Cursor, timestamp: datetime) -> None:
             cur.execute(f"DELETE FROM {self.table}")  # type: ignore
         cur.execute(
             f"""
-            INSERT INTO {self.table}({self.cols_str}, {",".join(self.meta_keys)})
-            SELECT {self.cols_str},{timestamp_placeholders} FROM {self.temp_table}
-            ON CONFLICT ({",".join(unique_key)}) {on_conflict}
+            INSERT INTO {self.table} AS t ({self.cols_str}, {", ".join(self.meta_keys)})
+            SELECT {self.cols_str}, {timestamp_placeholders} FROM {self.temp_table}
+            ON CONFLICT ({", ".join(unique_key)}) {on_conflict}
             """,  # type: ignore
             {"timestamp": timestamp},
+            binary=True,
         )
+        self.conn.commit()
+        self.upsert_timer.stop()
 
         if self.load_type == LoadType.INCREMENTAL and self.model.last_updated_date_table():
             key = self.model.last_updated_timestamp_col()
             last_updated_cols = self.model.last_updated_date_column()
             set_clause = ", ".join(f"{col} = %(timestamp)s" for col in last_updated_cols)
 
-            # We require multi-step transactions since we're dealing with temp tables, so there
-            # is a chance of a deadlock here.
-            # However, it's safe to ignore these because if the timestamp for this row is being
-            # updated concurrently then it's going to have the same end result anyway.
-            # If a deadlock occurs, the CTE returns no rows and this is a no-op.
-
-            cur.execute(
-                f"""
-                WITH current_ts AS (
-                    SELECT {key}
-                    FROM {self.model.last_updated_date_table()}
-                    WHERE {key} IN (
-                        SELECT {key} FROM {self.temp_table}
+            self.last_updated_timer.start()
+            try:
+                # We want to immediately terminate the transaction if there is already a lock on
+                # the table so that we avoid extraneous waits because if there is a lock this table
+                # is being updated concurrently and that existing update will have the same result
+                cur.execute("SAVEPOINT pre_last_updated")
+                cur.execute("SET LOCAL lock_timeout=1")
+                cur.execute("SET LOCAL statement_timeout=3000")
+                cur.execute(
+                    f"""
+                    WITH current_ts AS (
+                        SELECT {key}
+                        FROM {self.model.last_updated_date_table()}
+                        WHERE {key} IN (
+                            SELECT {key} FROM {self.temp_table}
+                        )
+                        ORDER BY {key}
+                        FOR UPDATE SKIP LOCKED
                     )
-                    ORDER BY {key}
-                    FOR UPDATE SKIP LOCKED
+                    UPDATE {self.model.last_updated_date_table()} u
+                    SET {set_clause}
+                    FROM current_ts t
+                    WHERE u.{key} = t.{key};
+                    """,  # type: ignore
+                    {"timestamp": timestamp},
                 )
-                UPDATE {self.model.last_updated_date_table()} u
-                SET {set_clause}
-                FROM current_ts t
-                WHERE u.{key} = t.{key};
-                """,  # type: ignore
-                {"timestamp": timestamp},
-            )
+                self.conn.commit()
+            except (DeadlockDetected, LockNotAvailable, QueryCanceled) as ex:
+                logger.warning(
+                    "deadlock/lock/statement timeout updating update timestamp, ignoring: %s", ex
+                )
+                cur.execute("ROLLBACK TO SAVEPOINT pre_last_updated")
+            self.last_updated_timer.stop()
 
     def _copy_data(self, cur: psycopg.Cursor, results: Sequence[T]) -> None:
         # Use COPY to load the batch into Postgres.

apps/bfd-pipeline-idr/test_pipeline.py

@@ -1,15 +1,13 @@
 import os
 import shutil
 import subprocess
-from collections.abc import Generator
 from datetime import datetime, timedelta
 from pathlib import Path
 from typing import cast
 from uuid import uuid4
 
 import psycopg
-import pytest
-from psycopg import sql
+from psycopg import Connection, sql
 from psycopg.rows import DictRow, dict_row
 from testcontainers.core.config import testcontainers_config  # type: ignore
 
@@ -45,7 +43,7 @@ def _run_migrator(postgres: PostgresContainer) -> None:
             f"-Dflyway.user={postgres.username} "
             f"-Dflyway.password={postgres.password} "
             "-Duser.timezone=UTC",
-            cwd="../bfd-db-migrator-ng",
+            cwd=Path(__file__).parent.joinpath("../bfd-db-migrator-ng"),
             shell=True,
             capture_output=True,
             check=True,
@@ -55,41 +53,7 @@ def _run_migrator(postgres: PostgresContainer) -> None:
         raise
 
 
-@pytest.fixture(scope="module")
-def setup_db() -> Generator[PostgresContainer]:
-    with PostgresContainer("postgres:16", driver="") as postgres:
-        with psycopg.connect(postgres.get_connection_url()) as conn:
-            with Path("./mock-idr.sql").open() as f:
-                conn.execute(f.read())  # type: ignore
-            conn.commit()
-
-            _run_migrator(postgres)
-            load_from_csv(PostgresExecutor(conn), "./test_samples1")  # type: ignore
-
-            info = conn.info
-            # Info level logs obscure the error output when running tests
-            # so we want to override this unless the calling process has set this explicitly
-            os.environ.setdefault("IDR_LOG_LEVEL", "warning")
-            os.environ["BFD_DB_ENDPOINT"] = info.host
-            os.environ["BFD_DB_PORT"] = str(info.port)
-            os.environ["BFD_DB_NAME"] = info.dbname
-            os.environ["BFD_DB_USERNAME"] = info.user
-            os.environ["BFD_DB_PASSWORD"] = info.password
-            os.environ["IDR_BATCH_SIZE"] = "100000"
-            os.environ["IDR_FORCE_LOAD_PROGRESS"] = "1"
-            os.environ["BFD_TEST_DATE"] = "2023-04-02"
-        yield postgres
-
-
-def test_pipeline(setup_db: PostgresContainer) -> None:
-    load_type = LoadType.INCREMENTAL
-    conn = cast(
-        psycopg.Connection[DictRow],
-        psycopg.connect(setup_db.get_connection_url(), row_factory=dict_row),  # type: ignore
-    )
-
-    conn.commit()
-
+def _do_test_pipeline(conn: Connection[DictRow], load_type: LoadType) -> None:
     run(Source.POSTGRES, LoadMode.SYNTHETIC, load_type)
 
     cur = conn.execute("select * from idr.beneficiary order by bene_sk")
@@ -521,3 +485,62 @@ def test_pipeline(setup_db: PostgresContainer) -> None:
 def advance_time(timestamp: datetime) -> None:
     new_time = timestamp + timedelta(minutes=1)
     os.environ["BFD_TEST_DATE"] = new_time.isoformat()
+
+
+def test_pipeline() -> None:
+    # This is REALLY dumb. Typically we would use a parameterized test plus a fixture to setup the
+    # database and run the same test with both load types, but GitHub Actions Runners seem to have
+    # some issue with testcontainers where only a single test case succeeds and all others fail to
+    # connect. This forces sequential execution of each test case and ensures only a single
+    # container is ever started, avoiding the issue in CI
+    # TODO: Don't do this, find a way to make parameterized tests work in CI
+    with (
+        PostgresContainer("postgres:16", driver="") as postgres,
+        # No idea why pyright is upset about "row_factory", but we need to tell it to ignore the
+        # argument type here.
+        psycopg.connect(conninfo=postgres.get_connection_url(), row_factory=dict_row) as conn,  # pyright: ignore[reportArgumentType]
+    ):
+        for load_type in [LoadType.INCREMENTAL, LoadType.INITIAL]:
+            # Truncate all tables first. See above for justification
+            conn.execute(
+                """
+                DO $$ DECLARE
+                    r RECORD;
+                BEGIN
+                    FOR r IN (SELECT tablename FROM pg_tables WHERE schemaname = 'idr') LOOP
+                        EXECUTE 'DROP TABLE idr.' || quote_ident(r.tablename) || ' CASCADE';
+                    END LOOP;
+
+                    FOR r IN (
+                        SELECT tablename FROM pg_tables WHERE schemaname = 'cms_vdm_view_mdcr_prd'
+                    ) LOOP
+                        EXECUTE 'DROP TABLE cms_vdm_view_mdcr_prd.'
+                            || quote_ident(r.tablename)
+                            || ' CASCADE';
+                    END LOOP;
+                END $$;
+                """
+            )
+            conn.commit()
+
+            with Path(__file__).parent.joinpath("./mock-idr.sql").open() as f:
+                conn.execute(f.read())  # type: ignore
+            conn.commit()
+
+            _run_migrator(postgres)
+            load_from_csv(PostgresExecutor(conn), Path(__file__).parent.joinpath("./test_samples1"))  # type: ignore
+
+            info = conn.info
+            # Info level logs obscure the error output when running tests
+            # so we want to override this unless the calling process has set this explicitly
+            os.environ.setdefault("IDR_LOG_LEVEL", "warning")
+            os.environ["BFD_DB_ENDPOINT"] = info.host
+            os.environ["BFD_DB_PORT"] = str(info.port)
+            os.environ["BFD_DB_NAME"] = info.dbname
+            os.environ["BFD_DB_USERNAME"] = info.user
+            os.environ["BFD_DB_PASSWORD"] = info.password
+            os.environ["IDR_BATCH_SIZE"] = "100000"
+            os.environ["IDR_FORCE_LOAD_PROGRESS"] = "1"
+            os.environ["BFD_TEST_DATE"] = "2023-04-02"
+
+            _do_test_pipeline(cast(Connection[DictRow], conn), load_type)

ops/services/04-idr-pipeline/lambda_src/run-idr-pipeline/app/main.py

@@ -95,7 +95,9 @@ class IdrContainerOverrides:
     @classmethod
     def from_invoke(cls, invoke_model: InvokeModel) -> IdrContainerOverrides:
         env_kvs = (
-            [EnvKeyValue(name=k, value=v) for k, v in invoke_model.env] if invoke_model.env else []
+            [EnvKeyValue(name=k, value=v) for k, v in invoke_model.env.items()]
+            if invoke_model.env
+            else []
         )
         command_seq = invoke_model.command.split() if invoke_model.command else None