[DEV-16003] Downmerge into QAT

usaspending_api/download/helpers/psql_helpers.py

@@ -5,10 +5,10 @@
 
 
 def build_psql_env(
-        dsn: str,
-        statement_timeout_hours: Optional[int] = None,
-        work_mem_mb: Optional[int] = None,
-        base_env: Optional[dict] = None
+    dsn: str,
+    statement_timeout_hours: Optional[int] = None,
+    work_mem_mb: Optional[int] = None,
+    base_env: Optional[dict] = None,
 ) -> dict:
     """Build PostgreSQL environment variables from a database connection string."""
 
@@ -24,7 +24,7 @@ def build_psql_env(
     env["PGPORT"] = str(db_url.port or 5432)
     env["PGUSER"] = db_url.username or "postgres"
     env["PGPASSWORD"] = db_url.password or ""
-    env["PGDATABASE"] = db_url.path.lstrip('/') if db_url.path else "postgres"
+    env["PGDATABASE"] = db_url.path.lstrip("/") if db_url.path else "postgres"
 
     # Set optional PostgreSQL options
     if statement_timeout_hours or work_mem_mb:
@@ -39,21 +39,18 @@ def build_psql_env(
 
 
 def run_psql_to_file(
-        sql_path: str,
-        output_path: str,
-        env: dict,
-        quiet: bool = True,
-        on_error_stop: bool = True
+    sql_path: str, output_path: str, env: dict, quiet: bool = True, on_error_stop: bool = True
 ) -> None:
     """
     Execute a psql command that reads SQL from a file and writes output to another file.
     """
     import logging
+
     logger = logging.getLogger(__name__)
 
     # Log the SQL file contents for debugging
     try:
-        with open(sql_path, 'r') as f:
+        with open(sql_path, "r") as f:
             sql_content = f.read()
             logger.info(f"SQL file contents (first 500 chars): {sql_content[:500]}")
     except Exception as e:
@@ -71,12 +68,7 @@ def run_psql_to_file(
 
     # Test database connection first
     logger.info("Testing database connection...")
-    test_process = subprocess.run(
-        ["psql", "-c", "SELECT 1;"],
-        env=env,
-        capture_output=True,
-        timeout=5
-    )
+    test_process = subprocess.run(["psql", "-c", "SELECT 1;"], env=env, capture_output=True, timeout=30)
     if test_process.returncode != 0:
         logger.error(f"Database connection test failed: {test_process.stderr.decode()}")
         raise Exception(f"Cannot connect to database: {test_process.stderr.decode()}")
@@ -103,13 +95,15 @@ def run_psql_to_file(
 
     # Wait for both processes to complete with timeout
     try:
-        psql_output, psql_error = psql_process.communicate(timeout=30)  # 30 second timeout
-        cat_process.wait(timeout=5)
+        psql_output, psql_error = psql_process.communicate()
+        cat_process.wait(timeout=30)
     except subprocess.TimeoutExpired:
         logger.error("Process timed out! Killing processes...")
         psql_process.kill()
         cat_process.kill()
-        raise Exception("psql process timed out after 30 seconds") from None
+        raise Exception(
+            "psql process timed out by the server's process OR cat process timed out after 30 seconds"
+        ) from None
 
     logger.info(f"psql return code: {psql_process.returncode}")
     logger.info(f"psql stdout: {psql_output.decode() if psql_output else 'empty'}")
@@ -119,11 +113,6 @@ def run_psql_to_file(
     if psql_process.returncode != 0:
         error_msg = psql_error.decode() if psql_error else psql_output.decode() if psql_output else "Unknown error"
         logger.error(f"psql failed: {error_msg}")
-        raise subprocess.CalledProcessError(
-            psql_process.returncode,
-            psql_args,
-            output=psql_output,
-            stderr=psql_error
-        )
+        raise subprocess.CalledProcessError(psql_process.returncode, psql_args, output=psql_output, stderr=psql_error)
 
     logger.info("psql completed successfully")