diff --git a/usaspending_api/download/helpers/psql_helpers.py b/usaspending_api/download/helpers/psql_helpers.py index db6a44f4e2..9ac71160fe 100644 --- a/usaspending_api/download/helpers/psql_helpers.py +++ b/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")