Addressing ruff issues

Author: zachflanders-frb

usaspending_api/etl/management/commands/archive_table_in_delta.py

@@ -1,18 +1,20 @@
 import logging
-import psycopg2
-
 from datetime import datetime, timedelta
-from django.core.management.base import BaseCommand
 
-from usaspending_api.common.helpers.sql_helpers import get_database_dsn_string
+import psycopg2
+from django.core.management.base import BaseCommand, CommandParser
+
 from usaspending_api.common.etl.spark import load_delta_table
 from usaspending_api.common.helpers.spark_helpers import (
     configure_spark_session,
     get_active_spark_session,
     get_jdbc_connection_properties,
     get_usas_jdbc_url,
 )
-from usaspending_api.download.delta_models.download_job import download_job_create_sql_string
+from usaspending_api.common.helpers.sql_helpers import get_database_dsn_string
+from usaspending_api.download.delta_models.download_job import (
+    download_job_create_sql_string,
+)
 from usaspending_api.etl.table_specs import ArchiveTableSpec
 
 logger = logging.getLogger(__name__)
@@ -38,7 +40,8 @@ class Command(BaseCommand):
     those records from Postgres.
     """
 
-    def add_arguments(self, parser):
+    @staticmethod
+    def add_arguments(parser: CommandParser) -> None:
         parser.add_argument(
             "--destination-table",
             type=str,
@@ -57,7 +60,8 @@ def add_arguments(self, parser):
             "--alt-db",
             type=str,
             required=False,
-            help="An alternate Delta Database (aka schema) in which to archive this table, overriding the TABLE_SPEC's destination_database",
+            help="An alternate Delta Database (aka schema) in which to archive this table, overriding the TABLE_SPEC's"
+            " destination_database",
         )
         parser.add_argument(
             "--alt-name",
@@ -66,7 +70,7 @@ def add_arguments(self, parser):
             help="An alternate Delta Table name which to archive this table, overriding the destination_table",
         )
 
-    def handle(self, *args, **options):
+    def handle(self, *args, **options) -> None:
         extra_conf = {
             # Config for Delta Lake tables and SQL. Need these to keep Dela table metadata in the metastore
             "spark.sql.extensions": "io.delta.sql.DeltaSparkSessionExtension",
@@ -107,14 +111,16 @@ def handle(self, *args, **options):
         # Resolve JDBC URL for Source Database
         jdbc_url = get_usas_jdbc_url()
         if not jdbc_url:
-            raise RuntimeError(f"Couldn't find JDBC url, please properly configure your CONFIG.")
+            raise RuntimeError(
+                "Couldn't find JDBC url, please properly configure your CONFIG."
+            )
         if not jdbc_url.startswith("jdbc:postgresql://"):
-            raise ValueError("JDBC URL given is not in postgres JDBC URL format (e.g. jdbc:postgresql://...")
+            raise ValueError(
+                "JDBC URL given is not in postgres JDBC URL format (e.g. jdbc:postgresql://..."
+            )
 
         # Retrieve data from Postgres
-        query_with_predicate = (
-            f"(SELECT * FROM {qualified_source_table} WHERE {archive_date_field} < '{archive_date_string}') AS tmp"
-        )
+        query_with_predicate = f"(SELECT * FROM {qualified_source_table} WHERE {archive_date_field} < '{archive_date_string}') AS tmp"
 
         df = spark.read.jdbc(
             url=jdbc_url,
@@ -125,7 +131,9 @@ def handle(self, *args, **options):
         # Write data to Delta Lake in Append Mode
         load_delta_table(spark, df, destination_table_name, overwrite=False)
         archived_count = df.count()
-        logger.info(f"Archived {archived_count} records from the {qualified_source_table}")
+        logger.info(
+            f"Archived {archived_count} records from the {qualified_source_table}"
+        )
 
         # Delete data from
         with psycopg2.connect(dsn=get_database_dsn_string()) as connection:
@@ -135,7 +143,9 @@ def handle(self, *args, **options):
                 )
                 deleted_count = cursor.rowcount
 
-        logger.info(f"Deleted {deleted_count} records from the {qualified_source_table} table")
+        logger.info(
+            f"Deleted {deleted_count} records from the {qualified_source_table} table"
+        )
 
         # Shut down spark
         if spark_created_by_command:

usaspending_api/etl/management/commands/create_delta_table.py

@@ -1,6 +1,6 @@
 import logging
 
-from django.core.management.base import BaseCommand
+from django.core.management.base import BaseCommand, CommandParser
 from pyspark.sql.types import StructType
 
 from usaspending_api.awards.delta_models.award_id_lookup import AWARD_ID_LOOKUP_SCHEMA
@@ -50,7 +50,7 @@ class Command(BaseCommand):
     This command creates an empty Delta Table based on the provided --destination-table argument.
     """
 
-    def add_arguments(self, parser):
+    def add_arguments(self, parser: CommandParser) -> None:
         parser.add_argument(
             "--destination-table",
             type=str,
@@ -79,7 +79,7 @@ def add_arguments(self, parser):
             "name",
         )
 
-    def handle(self, *args, **options):
+    def handle(self, *args, **options) -> None:
         spark = get_active_spark_session()
         spark_created_by_command = False
         if not spark:

usaspending_api/etl/management/commands/load_query_to_delta.py

@@ -2,7 +2,7 @@
 from argparse import ArgumentTypeError
 from typing import Callable
 
-from django.core.management.base import BaseCommand
+from django.core.management.base import BaseCommand, CommandParser
 from pyspark.sql import SparkSession
 
 from usaspending_api.common.etl.spark import create_ref_temp_views
@@ -35,7 +35,9 @@
     load_object_class_program_activity_incremental,
     object_class_program_activity_schema,
 )
-from usaspending_api.download.delta_models.transaction_download import transaction_download_schema
+from usaspending_api.download.delta_models.transaction_download import (
+    transaction_download_schema,
+)
 from usaspending_api.etl.table_specs import QueryTableSpec
 from usaspending_api.recipient.delta_models import (
     RECIPIENT_LOOKUP_POSTGRES_COLUMNS,
@@ -58,7 +60,10 @@
     AWARD_SEARCH_POSTGRES_GOLD_COLUMNS,
     award_search_create_sql_string,
 )
-from usaspending_api.search.delta_models.dataframes.award_search import load_award_search, load_award_search_incremental
+from usaspending_api.search.delta_models.dataframes.award_search import (
+    load_award_search,
+    load_award_search_incremental,
+)
 from usaspending_api.search.delta_models.dataframes.transaction_search import (
     load_transaction_search,
     load_transaction_search_incremental,
@@ -70,7 +75,12 @@
     subaward_search_create_sql_string,
     subaward_search_load_sql_string,
 )
-from usaspending_api.search.models import AwardSearch, SubawardSearch, SummaryStateView, TransactionSearch
+from usaspending_api.search.models import (
+    AwardSearch,
+    SubawardSearch,
+    SummaryStateView,
+    TransactionSearch,
+)
 from usaspending_api.settings import HOST
 from usaspending_api.transactions.delta_models import (
     SUMMARY_STATE_VIEW_COLUMNS,
@@ -226,8 +236,14 @@
                 "partition_keys": ["is_fpds"],
                 "partitioning_form": "LIST",
                 "partitions": [
-                    {"table_suffix": "_fpds", "partitioning_clause": "FOR VALUES IN (TRUE)"},
-                    {"table_suffix": "_fabs", "partitioning_clause": "FOR VALUES IN (FALSE)"},
+                    {
+                        "table_suffix": "_fpds",
+                        "partitioning_clause": "FOR VALUES IN (TRUE)",
+                    },
+                    {
+                        "table_suffix": "_fabs",
+                        "partitioning_clause": "FOR VALUES IN (FALSE)",
+                    },
                 ],
             },
         }
@@ -286,8 +302,11 @@
             "partition_column_type": "numeric",
             "delta_table_create_sql": account_balances_schema,
             "delta_table_create_options": {"delta.enableChangeDataFeed": True},
-            "column_names": list(),
-            "delta_table_create_partitions": ["reporting_fiscal_year", "funding_toptier_agency_id"],
+            "column_names": [],
+            "delta_table_create_partitions": [
+                "reporting_fiscal_year",
+                "funding_toptier_agency_id",
+            ],
         }
     ),
     "award_financial_download": QueryTableSpec(
@@ -299,8 +318,11 @@
             "partition_column_type": "numeric",
             "delta_table_create_sql": award_financial_schema,
             "delta_table_create_options": {"delta.enableChangeDataFeed": True},
-            "column_names": list(),
-            "delta_table_create_partitions": ["reporting_fiscal_year", "funding_toptier_agency_id"],
+            "column_names": [],
+            "delta_table_create_partitions": [
+                "reporting_fiscal_year",
+                "funding_toptier_agency_id",
+            ],
         }
     ),
     "object_class_program_activity_download": QueryTableSpec(
@@ -312,8 +334,11 @@
             "partition_column_type": "numeric",
             "delta_table_create_sql": object_class_program_activity_schema,
             "delta_table_create_options": {"delta.enableChangeDataFeed": True},
-            "column_names": list(),
-            "delta_table_create_partitions": ["reporting_fiscal_year", "funding_toptier_agency_id"],
+            "column_names": [],
+            "delta_table_create_partitions": [
+                "reporting_fiscal_year",
+                "funding_toptier_agency_id",
+            ],
         }
     ),
     "transaction_download": QueryTableSpec(
@@ -323,8 +348,12 @@
             "partition_column_type": "numeric",
             "delta_table_create_sql": transaction_download_schema,
             "delta_table_create_options": {"delta.enableChangeDataFeed": True},
-            "column_names": list(),
-            "delta_table_create_partitions": ["awarding_agency_code", "is_fpds", "action_date_fiscal_year"],
+            "column_names": [],
+            "delta_table_create_partitions": [
+                "awarding_agency_code",
+                "is_fpds",
+                "action_date_fiscal_year",
+            ],
         }
     ),
 }
@@ -342,7 +371,8 @@ class Command(BaseCommand):
     destination_table_name: str
     spark: SparkSession
 
-    def add_arguments(self, parser):
+    @staticmethod
+    def add_arguments(parser: CommandParser) -> None:
         parser.add_argument(
             "--destination-table",
             type=str,
@@ -370,7 +400,7 @@ def add_arguments(self, parser):
             help="Whether or not the table will be updated incrementally",
         )
 
-    def handle(self, *args, **options):
+    def handle(self, *args, **options) -> None:
         extra_conf = {
             # Config for Delta Lake tables and SQL. Need these to keep Dela table metadata in the metastore
             "spark.sql.extensions": "io.delta.sql.DeltaSparkSessionExtension",
@@ -385,17 +415,25 @@ def handle(self, *args, **options):
         spark_created_by_command = False
         if not self.spark:
             spark_created_by_command = True
-            self.spark = configure_spark_session(**extra_conf, spark_context=self.spark)  # type: SparkSession
+            self.spark = configure_spark_session(
+                **extra_conf, spark_context=self.spark
+            )  # type: SparkSession
 
         # Resolve Parameters
         destination_table = options["destination_table"]
         table_spec = TABLE_SPEC[destination_table]
         self.destination_database = options["alt_db"] or table_spec.destination_database
-        self.destination_table_name = options["alt_name"] or destination_table.split(".")[-1]
-        source_query_key = "source_query_incremental" if options["incremental"] else "source_query"
+        self.destination_table_name = (
+            options["alt_name"] or destination_table.split(".")[-1]
+        )
+        source_query_key = (
+            "source_query_incremental" if options["incremental"] else "source_query"
+        )
         load_query = getattr(table_spec, source_query_key)
         if load_query is None:
-            raise ArgumentTypeError(f"Invalid source query. `{source_query_key}` must be specified in the TABLE_SPEC.")
+            raise ArgumentTypeError(
+                f"Invalid source query. `{source_query_key}` must be specified in the TABLE_SPEC."
+            )
 
         # Set the database that will be interacted with for all Delta Lake table Spark-based activity
         logger.info(f"Using Spark Database: {self.destination_database}")
@@ -405,15 +443,19 @@ def handle(self, *args, **options):
 
         if isinstance(load_query, list):
             for index, query in enumerate(load_query):
-                logger.info(f"Running query number: {index + 1}\nPreview of query: {query[:100]}")
+                logger.info(
+                    f"Running query number: {index + 1}\nPreview of query: {query[:100]}"
+                )
                 self.run_spark_sql(query)
         else:
             self.run_spark_sql(load_query)
 
         if spark_created_by_command:
             self.spark.stop()
 
-    def run_spark_sql(self, query: str | Callable[[SparkSession, str, str], None]):
+    def run_spark_sql(
+        self, query: str | Callable[[SparkSession, str, str], None]
+    ) -> None:
         if isinstance(query, str):
             jdbc_conn_props = get_jdbc_connection_properties()
             self.spark.sql(
@@ -430,4 +472,6 @@ def run_spark_sql(self, query: str | Callable[[SparkSession, str, str], None]):
         elif isinstance(query, Callable):
             query(self.spark, self.destination_database, self.destination_table_name)
         else:
-            raise ArgumentTypeError(f"Invalid query. `{query}` must be a string or a Callable.")
+            raise ArgumentTypeError(
+                f"Invalid query. `{query}` must be a string or a Callable."
+            )