BFD-4192: Configure additional CI checks for Python projects (#2758)

Author: aschey-forpeople

.github/workflows/ci-pytest.yml .github/workflows/ci-python.yml.github/workflows/ci-pytest.yml renamed to .github/workflows/ci-python.yml

@@ -6,6 +6,7 @@ on:
       - "**/*.py"
       - "**/*requirement*.txt"
       - "**/uv.lock"
+      - ".github/workflows/ci-python.yml"
 
   merge_group:
 
@@ -24,8 +25,27 @@ jobs:
               projectName: "sftp-outbound-transfer-lambda",
               projectRoot: "ops/terraform/services/eft/lambda_src/sftp_outbound_transfer",
               testsRoot: "ops/terraform/services/eft/lambda_src/sftp_outbound_transfer/tests",
+              uvCommand: "uv sync",
+              pytestCommand: "uv run pytest",
               pythonVersion: "3.13",
             },
+            {
+              projectName: "locust",
+              projectRoot: "apps/utils/locust_tests",
+              testsRoot: "apps/utils/locust_tests",
+              uvCommand: "uv sync --group lambda-run-locust",
+              # no tests
+              pytestCommand: "echo 'no tests'",
+              pythonVersion: "3.13"
+            },
+            {
+              projectName: "idr-pipeline",
+              projectRoot: "apps/bfd-pipeline/bfd-pipeline-idr",
+              testsRoot: "apps/bfd-pipeline/bfd-pipeline-idr",
+              uvCommand: "uv sync",
+              pytestCommand: "uv run pytest",
+              pythonVersion: "3.13"
+            }
           ]
       fail-fast: false
     steps:
@@ -45,5 +65,12 @@ jobs:
           echo "Testing ${{ matrix.projects.projectName }}"
           cd ${{ github.workspace }}/${{ matrix.projects.projectRoot }}
           uv python install
-          uv sync
-          uv run pytest ${{ github.workspace }}/${{ matrix.projects.testsRoot }}
+          ${{ matrix.projects.uvCommand }}
+          # check linting
+          uv run ruff check
+          # check formatting
+          uv run ruff format --check
+          # check types
+          uv run pyright .
+          # tests
+          ${{ matrix.projects.pytestCommand }} ${{ github.workspace }}/${{ matrix.projects.testsRoot }}

apps/bfd-pipeline/bfd-pipeline-idr/extractor.py

@@ -79,7 +79,8 @@ def extract_idr_data(self, cls: type[T], progress: LoadProgress | None) -> Itera
                 "{WHERE_CLAUSE}",
                 f"""
                     WHERE 
-                        ({update_timestamp_col} IS NOT NULL AND {batch_timestamp_col} {op} %(timestamp)s {update_clause})
+                        ({update_timestamp_col} IS NOT NULL 
+                            AND {batch_timestamp_col} {op} %(timestamp)s {update_clause})
                         AND {batch_timestamp_col} {op} '{get_min_transaction_date()}' 
                     """,
             ).replace("{ORDER_BY}", f"ORDER BY {batch_timestamp_col}"),
@@ -133,7 +134,8 @@ def extract_many(self, cls: type[T], sql: str, params: dict[str, DbType]) -> Ite
             cursor_execute_timer.stop()
 
             cursor_fetch_timer.start()
-            # fetchmany can return list[dict] or list[tuple] but we'll only use queries that return dicts
+            # fetchmany can return list[dict] or list[tuple] but we'll only use
+            # queries that return dicts
             batch: list[dict[str, DbType]] = cur.fetchmany(self.batch_size)  # type: ignore[assignment]
             cursor_fetch_timer.stop()
 

apps/bfd-pipeline/bfd-pipeline-idr/loader.py

@@ -24,7 +24,8 @@ def print_timers() -> None:
 def get_connection_string() -> str:
     port = os.environ.get("BFD_DB_PORT") or "5432"
     dbname = os.environ.get("BFD_DB_NAME") or "idr"
-    return f"host={os.environ['BFD_DB_ENDPOINT']} port={port} dbname={dbname} user={os.environ['BFD_DB_USERNAME']} password={os.environ['BFD_DB_PASSWORD']}"
+    return f"host={os.environ['BFD_DB_ENDPOINT']} port={port} dbname={dbname} \
+        user={os.environ['BFD_DB_USERNAME']} password={os.environ['BFD_DB_PASSWORD']}"
 
 
 class PostgresLoader:
@@ -52,17 +53,21 @@ def load(
         update_set = ", ".join([f"{v}=EXCLUDED.{v}" for v in insert_cols if v not in unique_key])
         timestamp = datetime.now(UTC)
         table = model.table()
-        # trim the schema from the table name to create the temp table (temp tables can't be created with an explicit schema set)
+        # trim the schema from the table name to create the temp table
+        # (temp tables can't be created with an explicit schema set)
         temp_table = table.split(".")[1] + "_temp"
         with self.conn.cursor() as cur:
             # load each batch in a separate transaction
             for results in fetch_results:
                 # Load each batch into a temp table
-                # This is necessary because we want to use COPY to quickly transfer everything into Postgres
-                # but COPY can't handle constraint conflicts natively.
+                # This is necessary because we want to use COPY to quickly
+                # transfer everything into Postgres, but COPY can't handle
+                # constraint conflicts natively.
+                #
                 # Note that temp tables don't use WAL so that helps with throughput as well.
                 #
-                # For simplicity's sake, we'll create our temp tables using the existing schema and just drop the columns we need to ignore
+                # For simplicity's sake, we'll create our temp tables using the existing schema and
+                # just drop the columns we need to ignore.
                 temp_table_timer.start()
                 cur.execute(
                     f"CREATE TEMPORARY TABLE {temp_table} (LIKE {table}) ON COMMIT DROP"  # type: ignore
@@ -78,12 +83,13 @@ def load(
                 temp_table_timer.stop()
 
                 # Use COPY to load the batch into Postgres.
-                # COPY has a number of optimizations that make bulk loading more efficient than a bunch of INSERTs.
-                # The entire operation is performed in a single statement, resulting in fewer network round-trips,
-                # less WAL activity, and less context switching.
+                # COPY has a number of optimizations that make bulk loading more efficient
+                # than a bunch of INSERTs.
+                # The entire operation is performed in a single statement, resulting in
+                # fewer network round-trips, less WAL activity, and less context switching.
 
-                # Even though we need to move the data from the temp table in the next step, it should still be
-                # faster than alternatives.
+                # Even though we need to move the data from the temp table in the next step,
+                # it should still be faster than alternatives.
                 copy_timer.start()
                 with cur.copy(f"COPY {temp_table} ({cols_str}) FROM STDIN") as copy:  # type: ignore
                     for row in results:
@@ -92,8 +98,10 @@ def load(
                 copy_timer.stop()
 
                 if len(results) > 0:
-                    # For immutable tables, we may still be attempting to re-load some data due to a batch cancellation.
-                    # In these cases, we can assume any conflicting rows have already been loaded so "DO NOTHING" is appropriate here.
+                    # For immutable tables, we may still be attempting to re-load some data
+                    # due to a batch cancellation.
+                    # In these cases, we can assume any conflicting rows have already been loaded so
+                    # "DO NOTHING" is appropriate here.
                     on_conflict = (
                         "DO NOTHING"
                         if immutable
@@ -113,7 +121,8 @@ def load(
                     insert_timer.stop()
 
                     last = results[len(results) - 1].model_dump()
-                    # Some tables that contain reference data (like contract info) may not have the normal IDR timestamps
+                    # Some tables that contain reference data (like contract info) may not have the
+                    # normal IDR timestamps.
                     # For now we won't support incremental refreshes for those tables
                     batch_timestamp_col = model.batch_timestamp_col(
                         progress is None or progress.is_historical()

apps/bfd-pipeline/bfd-pipeline-idr/model.py

@@ -1,6 +1,6 @@
 from collections.abc import Iterable
 from datetime import UTC, date, datetime
-from typing import Annotated, Optional, TypeVar
+from typing import Annotated, TypeVar
 
 from pydantic import BaseModel, BeforeValidator
 
@@ -92,7 +92,7 @@ def batch_timestamp_col(cls, is_historical: bool) -> str | None:
         return cls._get_timestamp_col(BATCH_TIMESTAMP)
 
     @classmethod
-    def batch_timestamp_col_alias(cls, is_historical: bool) -> Optional[str]:
+    def batch_timestamp_col_alias(cls, is_historical: bool) -> str | None:
         col = cls.batch_timestamp_col(is_historical)
         if col:
             return cls._format_column_alias(col)
@@ -377,7 +377,7 @@ class IdrBeneficiaryXref(IdrBaseModel):
     src_rec_ctre_ts: Annotated[datetime, {PRIMARY_KEY: True, BATCH_TIMESTAMP: True}]
 
     @staticmethod
-    def table():
+    def table() -> str:
         return "idr.beneficiary_xref"
 
     @staticmethod
@@ -408,10 +408,12 @@ def table() -> str:
 
     @staticmethod
     def _current_fetch_query() -> str:
-        # equivalent to "select distinct on", but Snowflake has different syntax for that so it's unfortunately not portable
+        # equivalent to "select distinct on", but Snowflake has different syntax for that,
+        # so it's unfortunately not portable
         return """
             WITH dupes as (
-                SELECT {COLUMNS}, ROW_NUMBER() OVER (PARTITION BY bene_sk, cntrct_pbp_sk, bene_enrlmt_efctv_dt 
+                SELECT {COLUMNS}, ROW_NUMBER() OVER (
+                    PARTITION BY bene_sk, cntrct_pbp_sk, bene_enrlmt_efctv_dt 
                 {ORDER_BY} DESC) as row_order
                 FROM cms_vdm_view_mdcr_prd.v2_mdcr_bene_elctn_prd_usg
                 {WHERE_CLAUSE}
@@ -722,7 +724,8 @@ def table() -> str:
     def _current_fetch_query() -> str:
         clm = ALIAS_CLM
         line = ALIAS_LINE
-        # Note: joining on clm_uniq_id isn't strictly necessary, but it did seem to improve performance a bit
+        # Note: joining on clm_uniq_id isn't strictly necessary,
+        # but it did seem to improve performance a bit
         return f"""
             SELECT {{COLUMNS}}
             FROM cms_vdm_view_mdcr_prd.v2_mdcr_clm {clm}
@@ -741,8 +744,7 @@ def _current_fetch_query() -> str:
 def transform_default_hipps_code(value: str | None) -> str:
     if value is None or value == "00000":
         return ""
-    else:
-        return value
+    return value
 
 
 class IdrClaimLineInstitutional(IdrBaseModel):

apps/bfd-pipeline/bfd-pipeline-idr/pyproject.toml

@@ -11,7 +11,13 @@ dependencies = [
 ]
 
 [dependency-groups]
-dev = ["ruff>=0.11.11", "uv>=0.7.8", "pytest>=8.3.5", "testcontainers>=4.9.2"]
+dev = [
+    "ruff>=0.11.11",
+    "uv>=0.7.8",
+    "pytest>=8.3.5",
+    "testcontainers>=4.9.2",
+    "pyright>=1.1.403",
+]
 
 # [[tool.mypy.overrides]]
 # module = ["testcontainers.*"]
@@ -53,6 +59,8 @@ select = [
     "FURB",
     "RUF",
 ]
+# Don't require docstrings
+ignore = ["D100", "D101", "D102", "D103", "D104", "D105", "D107"]
 
 [tool.ruff.lint.pydocstyle]
 convention = "pep257"

apps/bfd-pipeline/bfd-pipeline-idr/test_pipeline.py

@@ -1,5 +1,7 @@
 import time
-from datetime import datetime, timezone, UTC
+from collections.abc import Generator
+from datetime import UTC, datetime
+from pathlib import Path
 from typing import cast
 
 import psycopg
@@ -20,14 +22,16 @@
 
 
 @pytest.fixture(scope="session", autouse=True)
-def psql_url():
+def psql_url() -> Generator[str]:
     with PostgresContainer("postgres:16", driver="") as postgres:
         psql_url = postgres.get_connection_url()
         conn = psycopg.connect(psql_url)
 
-        conn.execute(open("./mock-idr.sql", "r").read())  # type: ignore
+        with Path("./mock-idr.sql").open() as f:
+            conn.execute(f.read())  # type: ignore
         conn.commit()
-        conn.execute(open("./bfd.sql", "r").read())  # type: ignore
+        with Path("./bfd.sql").open() as f:
+            conn.execute(f.read())  # type: ignore
         conn.commit()
 
         load_from_csv(conn, "./test_samples1")
@@ -36,7 +40,7 @@ def psql_url():
 
 
 class TestPipeline:
-    def test_pipeline(self, psql_url: str):
+    def test_pipeline(self, psql_url: str) -> None:
         run_pipeline(PostgresExtractor(psql_url, 100_000), psql_url)
         conn = cast(psycopg.Connection[DictRow], psycopg.connect(psql_url, row_factory=dict_row))  # type: ignore
         cur = conn.execute("select * from idr.beneficiary order by bene_sk")

apps/bfd-pipeline/bfd-pipeline-idr/timer.py

@@ -2,16 +2,16 @@
 
 
 class Timer:
-    def __init__(self, name: str):
+    def __init__(self, name: str) -> None:
         self.elapsed = 0.0
         self.perf_start = 0.0
         self.name = name
 
-    def start(self):
+    def start(self) -> None:
         self.perf_start = time.perf_counter()
 
-    def stop(self):
+    def stop(self) -> None:
         self.elapsed += time.perf_counter() - self.perf_start
 
-    def print_results(self):
+    def print_results(self) -> None:
         print(f"Time taken for {self.name}: {self.elapsed:.6f} seconds")