From e598535d4e2a526c9efc41d2478b6f5a8bdeff2e Mon Sep 17 00:00:00 2001 From: Lee Chagolla-Christensen Date: Sat, 25 Oct 2025 17:54:19 -0700 Subject: [PATCH] more psycopg to asyncpg pools in file enrichment - Replace psycopg_pool.ConnectionPool with asyncpg.Pool throughout file_enrichment service - Store asyncpg pool in FastAPI app.state instead of global variable - Update type hints - Fix enrichment success tracking query typo in noseyparker - Improve formatting in dotnet finding summaries (use code formatting for method names) --- .../file_enrichment/controller.py | 14 +- .../file_enrichment/file_enrichment/dotnet.py | 152 ++++++++---------- .../file_enrichment/noseyparker.py | 117 +++++++------- .../file_enrichment/routes/enrichments.py | 114 ++++++------- .../subscriptions/dotnet_handler.py | 6 +- .../subscriptions/file_handler.py | 91 +++++------ .../subscriptions/noseyparker_handler.py | 4 +- .../file_enrichment/workflow_recovery.py | 108 ++++++------- 8 files changed, 276 insertions(+), 330 deletions(-) diff --git a/projects/file_enrichment/file_enrichment/controller.py b/projects/file_enrichment/file_enrichment/controller.py index 84b7c89..238057f 100644 --- a/projects/file_enrichment/file_enrichment/controller.py +++ b/projects/file_enrichment/file_enrichment/controller.py @@ -14,7 +14,6 @@ from file_enrichment.postgres_notifications import postgres_notify_listener from file_enrichment.workflow_recovery import recover_interrupted_workflows from nemesis_dpapi import DpapiManager as NemesisDpapiManager from nemesis_dpapi.eventing import DaprDpapiEventPublisher -from psycopg_pool import ConnectionPool from .debug_utils import setup_debug_signals from .routes.dpapi import dpapi_background_monitor, dpapi_router @@ -34,8 +33,6 @@ max_workflow_execution_time = int( logger.info(f"max_workflow_execution_time: {max_workflow_execution_time}", pid=os.getpid()) -pool = ConnectionPool(get_postgres_connection_str(), open=True) - module_execution_order = [] workflow_manager: WorkflowManager = None @@ -60,7 +57,6 @@ async def lifespan(app: FastAPI): app.state.event_loop = asyncio.get_running_loop() set_fastapi_loop(asyncio.get_event_loop()) - # Create asyncpg connection pool for WorkflowManager and workflow activities dapr_client = DaprClient() postgres_connection_string = get_postgres_connection_str(dapr_client) @@ -69,7 +65,7 @@ async def lifespan(app: FastAPI): min_size=5, max_size=15, ) - logger.info("AsyncPG pool created", pid=os.getpid()) + app.state.asyncpg_pool = asyncpg_pool # Initialize global DpapiManager for the application lifetime dpapi_manager = NemesisDpapiManager( @@ -107,7 +103,7 @@ async def lifespan(app: FastAPI): logger.info("Started masterkey watcher task", pid=os.getpid()) # Recover any interrupted workflows before starting normal processing - await recover_interrupted_workflows(pool) + await recover_interrupted_workflows(asyncpg_pool) logger.info( "Workflow runtime initialized successfully", @@ -249,19 +245,19 @@ async def debug_tasks(): async def process_file(event: CloudEvent[File]): """Handler for incoming file events""" global workflow_manager - await process_file_event(event.data, workflow_manager, module_execution_order, pool) + await process_file_event(event.data, workflow_manager, module_execution_order, app.state.asyncpg_pool) @dapr_app.subscribe(pubsub="pubsub", topic="dotnet-output") async def process_dotnet_results(event: CloudEvent): """Handler for incoming .NET processing results from the dotnet_service.""" - await process_dotnet_event(event.data, pool) + await process_dotnet_event(event.data, app.state.asyncpg_pool) @dapr_app.subscribe(pubsub="pubsub", topic="noseyparker-output") async def process_nosey_parker_results(event: CloudEvent): """Handler for incoming Nosey Parker scan results""" - await process_noseyparker_event(event.data, pool) + await process_noseyparker_event(event.data, app.state.asyncpg_pool) @dapr_app.subscribe(pubsub="pubsub", topic="bulk-enrichment-task") diff --git a/projects/file_enrichment/file_enrichment/dotnet.py b/projects/file_enrichment/file_enrichment/dotnet.py index 49e3f3b..b71e527 100644 --- a/projects/file_enrichment/file_enrichment/dotnet.py +++ b/projects/file_enrichment/file_enrichment/dotnet.py @@ -1,16 +1,16 @@ # src/workflow/dotnet.py -import asyncio import json import uuid from typing import Any -import psycopg +import asyncpg from common.helpers import sanitize_for_jsonb from common.logger import get_logger from common.models import ( DotNetAssemblyAnalysis, EnrichmentResult, File, + FileEnriched, FileObject, Finding, FindingCategory, @@ -18,7 +18,6 @@ from common.models import ( Transform, ) from dapr.clients import DaprClient -from psycopg_pool import ConnectionPool logger = get_logger(__name__) @@ -59,9 +58,9 @@ def create_dotnet_finding_summary(analysis: DotNetAssemblyAnalysis) -> str: summary_content += f"\n#### {category}\n" for method in methods: if hasattr(method, "MethodName"): - summary_content += f"* **{method.MethodName}** (Level: {method.FilterLevel})\n" + summary_content += f"* `{method.MethodName}` (Level: {method.FilterLevel})\n" else: - summary_content += f"* {method}\n" + summary_content += f"* `{method}`\n" summary_content += "\n" return summary_content return "" @@ -78,10 +77,10 @@ def create_dotnet_finding_summary(analysis: DotNetAssemblyAnalysis) -> str: async def store_dotnet_results( object_id: str, - decompilation_object_id: str = None, - analysis: DotNetAssemblyAnalysis = None, - pool: ConnectionPool = None, - file_enriched=None, + decompilation_object_id: str | None = None, + analysis: DotNetAssemblyAnalysis | None = None, + pool: asyncpg.Pool | None = None, + file_enriched: FileEnriched | None = None, ): """ Store DotNet analysis results in the database, including creating findings and transforms. @@ -90,23 +89,22 @@ async def store_dotnet_results( object_id (str): The object ID of the file that was analyzed decompilation_object_id (str, optional): Object ID of the decompiled source ZIP analysis (DotNetAssemblyAnalysis, optional): Assembly analysis results - pool (ConnectionPool, optional): Database connection pool + pool (asyncpg.Pool, optional): Database connection pool file_enriched: The FileEnriched object for the original file """ try: # Update workflow success status try: - with pool.connection() as conn: - with conn.cursor() as cur: - cur.execute( - """ - UPDATE workflows - SET enrichments_success = array_append(enrichments_success, %s) - WHERE object_id = %s - """, - ("dotnet_service", object_id), - ) - conn.commit() + async with pool.acquire() as conn: + await conn.execute( + """ + UPDATE workflows + SET enrichments_success = array_append(enrichments_success, $1) + WHERE object_id = $2 + """, + "dotnet_service", + object_id, + ) except Exception as db_error: logger.error(f"Failed to update dotnet_service enrichment success in database: {str(db_error)}") @@ -203,71 +201,63 @@ async def store_dotnet_results( # Add findings to enrichment result enrichment_result.findings = findings_list - def store_in_db(): - with pool.connection() as conn: - with conn.cursor() as cur: - # Store main enrichment result - results_escaped = json.dumps(sanitize_for_jsonb(enrichment_result.model_dump(mode="json"))) - cur.execute( + # Store in database + async with pool.acquire() as conn: + # Store main enrichment result + results_escaped = json.dumps(sanitize_for_jsonb(enrichment_result.model_dump(mode="json"))) + await conn.execute( + """ + INSERT INTO enrichments (object_id, module_name, result_data) + VALUES ($1, $2, $3) + """, + object_id, + "dotnet_service", + results_escaped, + ) + + # Store any transforms + if enrichment_result.transforms: + for transform in enrichment_result.transforms: + await conn.execute( """ - INSERT INTO enrichments (object_id, module_name, result_data) - VALUES (%s, %s, %s) + INSERT INTO transforms (object_id, type, transform_object_id, metadata) + VALUES ($1, $2, $3, $4) """, - (object_id, "dotnet_service", results_escaped), + object_id, + transform.type, + transform.object_id, + json.dumps(transform.metadata) if transform.metadata else None, ) - # Store any transforms - if enrichment_result.transforms: - for transform in enrichment_result.transforms: - cur.execute( - """ - INSERT INTO transforms (object_id, type, transform_object_id, metadata) - VALUES (%s, %s, %s, %s) - """, - ( - object_id, - transform.type, - transform.object_id, - json.dumps(transform.metadata) if transform.metadata else None, - ), - ) + # Store any findings + for finding in findings_list: + # Convert each FileObject to a JSON string + data_as_strings = [] + for obj in finding.data: + # Convert the model to a dict first + if hasattr(obj, "model_dump"): + obj_dict = obj.model_dump() + else: + obj_dict = obj + sanitized_obj = sanitize_for_jsonb(obj_dict) + data_as_strings.append(json.dumps(sanitized_obj)) - # Store any findings - for finding in findings_list: - # Convert each FileObject to a JSON string - data_as_strings = [] - for obj in finding.data: - # Convert the model to a dict first - if hasattr(obj, "model_dump"): - obj_dict = obj.model_dump() - else: - obj_dict = obj - sanitized_obj = sanitize_for_jsonb(obj_dict) - data_as_strings.append(json.dumps(sanitized_obj)) - - cur.execute( - """ - INSERT INTO findings ( - finding_name, category, severity, object_id, - origin_type, origin_name, raw_data, data - ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s) - """, - ( - finding.finding_name, - finding.category, - finding.severity, - object_id, - finding.origin_type, - finding.origin_name, - json.dumps(sanitize_for_jsonb(finding.raw_data)), - json.dumps(data_as_strings), # Store as array of JSON strings - ), - ) - - conn.commit() - - # Run database operations in thread - await asyncio.to_thread(store_in_db) + await conn.execute( + """ + INSERT INTO findings ( + finding_name, category, severity, object_id, + origin_type, origin_name, raw_data, data + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + """, + finding.finding_name, + finding.category, + finding.severity, + object_id, + finding.origin_type, + finding.origin_name, + json.dumps(sanitize_for_jsonb(finding.raw_data)), + json.dumps(data_as_strings), # Store as array of JSON strings + ) logger.info("Successfully stored DotNet results", object_id=object_id, has_findings=len(findings_list) > 0) diff --git a/projects/file_enrichment/file_enrichment/noseyparker.py b/projects/file_enrichment/file_enrichment/noseyparker.py index 427ee13..111e4a8 100644 --- a/projects/file_enrichment/file_enrichment/noseyparker.py +++ b/projects/file_enrichment/file_enrichment/noseyparker.py @@ -1,4 +1,3 @@ -import asyncio import base64 import json import re @@ -7,11 +6,10 @@ import uuid from datetime import datetime from typing import Any -import psycopg +import asyncpg from common.helpers import sanitize_for_jsonb from common.logger import get_logger from common.models import EnrichmentResult, FileObject, Finding, FindingCategory, FindingOrigin, MatchInfo, ScanStats -from psycopg_pool import ConnectionPool logger = get_logger(__name__) @@ -162,7 +160,7 @@ async def store_noseyparker_results( object_id: str, matches: list[MatchInfo], scan_stats: ScanStats, - pool: ConnectionPool, + pool: asyncpg.Pool, ): """ Store Nosey Parker results in the database, including creating findings. @@ -171,21 +169,20 @@ async def store_noseyparker_results( object_id (str): The object ID of the file that was scanned matches (List[MatchInfo]): List of match information from Nosey Parker scan_stats (dict, optional): Statistics about the scan - pool (ConnectionPool): Database connection pool + pool (asyncpg.Pool): Database connection pool """ try: try: - with pool.connection() as conn: - with conn.cursor() as cur: - cur.execute( - """ - UPDATE workflows - SET enrichments_success = array_append(enrichments_failure, %s) - WHERE object_id = %s - """, - ("noseyparker", object_id), - ) - conn.commit() + async with pool.acquire() as conn: + await conn.execute( + """ + UPDATE workflows + SET enrichments_success = array_append(enrichments_failure, $1) + WHERE object_id = $2 + """, + "noseyparker", + object_id, + ) except Exception as db_error: logger.error(f"Failed to update noseyparker enrichment success in database: {str(db_error)}") @@ -233,55 +230,49 @@ async def store_noseyparker_results( # Add findings to enrichment result enrichment_result.findings = findings_list - def store_in_db(): - with pool.connection() as conn: - with conn.cursor() as cur: - # Store main enrichment result - results_escaped = json.dumps(sanitize_for_jsonb(enrichment_result.model_dump(mode="json"))) - cur.execute( - """ - INSERT INTO enrichments (object_id, module_name, result_data) - VALUES (%s, %s, %s) - """, - (object_id, "noseyparker", results_escaped), - ) + # Store in database + async with pool.acquire() as conn: + # Store main enrichment result + results_escaped = json.dumps(sanitize_for_jsonb(enrichment_result.model_dump(mode="json"))) + await conn.execute( + """ + INSERT INTO enrichments (object_id, module_name, result_data) + VALUES ($1, $2, $3) + """, + object_id, + "noseyparker", + results_escaped, + ) - # Store any findings - for finding in findings_list: - # Convert each FileObject to a JSON string - data_as_strings = [] - for obj in finding.data: - # Convert the model to a dict first - if hasattr(obj, "model_dump"): - obj_dict = obj.model_dump() - else: - obj_dict = obj - sanitized_obj = sanitize_for_jsonb(obj_dict) - data_as_strings.append(json.dumps(sanitized_obj)) + # Store any findings + for finding in findings_list: + # Convert each FileObject to a JSON string + data_as_strings = [] + for obj in finding.data: + # Convert the model to a dict first + if hasattr(obj, "model_dump"): + obj_dict = obj.model_dump() + else: + obj_dict = obj + sanitized_obj = sanitize_for_jsonb(obj_dict) + data_as_strings.append(json.dumps(sanitized_obj)) - cur.execute( - """ - INSERT INTO findings ( - finding_name, category, severity, object_id, - origin_type, origin_name, raw_data, data - ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s) - """, - ( - finding.finding_name, - finding.category, - finding.severity, - object_id, - finding.origin_type, - finding.origin_name, - json.dumps(sanitize_for_jsonb(finding.raw_data)), - json.dumps(data_as_strings), # Store as array of JSON strings - ), - ) - - conn.commit() - - # Run database operations in thread - await asyncio.to_thread(store_in_db) + await conn.execute( + """ + INSERT INTO findings ( + finding_name, category, severity, object_id, + origin_type, origin_name, raw_data, data + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + """, + finding.finding_name, + finding.category, + finding.severity, + object_id, + finding.origin_type, + finding.origin_name, + json.dumps(sanitize_for_jsonb(finding.raw_data)), + json.dumps(data_as_strings), # Store as array of JSON strings + ) logger.info("Successfully stored NoseyParker results", object_id=object_id, match_count=len(matches)) diff --git a/projects/file_enrichment/file_enrichment/routes/enrichments.py b/projects/file_enrichment/file_enrichment/routes/enrichments.py index 71fe510..96400e6 100644 --- a/projects/file_enrichment/file_enrichment/routes/enrichments.py +++ b/projects/file_enrichment/file_enrichment/routes/enrichments.py @@ -5,9 +5,10 @@ import json import os import uuid +import asyncpg import common.helpers as helpers from common.logger import get_logger -from fastapi import APIRouter, Body, HTTPException, Path +from fastapi import APIRouter, Body, HTTPException, Path, Request from file_enrichment.workflow import wf_runtime from pydantic import BaseModel @@ -60,13 +61,13 @@ async def list_enrichments(): @router.post("/enrichments/{enrichment_name}") async def run_enrichment( enrichment_name: str = Path(..., description="Name of the enrichment module to run"), - request: EnrichmentRequest = Body(..., description="The enrichment request containing the object ID"), + enrichment_request: EnrichmentRequest = Body(..., description="The enrichment request containing the object ID"), + request: Request = None, ): """Run a specific enrichment module directly.""" - # Import pool from controller module to avoid circular imports - from file_enrichment import controller - try: + # Get asyncpg pool from app state + asyncpg_pool: asyncpg.Pool = request.app.state.asyncpg_pool # Check if module if not wf_runtime or not wf_runtime.modules: raise HTTPException(status_code=503, detail="Workflow runtime or modules not initialized") @@ -78,80 +79,71 @@ async def run_enrichment( module = wf_runtime.modules[enrichment_name] # Check if we should process this file - run in thread since it might use sync operations - should_process = await asyncio.to_thread(module.should_process, request.object_id) + should_process = await asyncio.to_thread(module.should_process, enrichment_request.object_id) if not should_process: return { "status": "skipped", "message": f"Module {enrichment_name} decided to skip processing", - "object_id": request.object_id, + "object_id": enrichment_request.object_id, "instance_id": "", } # Process the file in a separate thread to avoid event loop conflicts - result = await asyncio.to_thread(module.process, request.object_id) + result = await asyncio.to_thread(module.process, enrichment_request.object_id) if result: # Store enrichment result in database - def store_results(): - with controller.pool.connection() as conn: - with conn.cursor() as cur: - # Store main enrichment result - results_escaped = json.dumps(helpers.sanitize_for_jsonb(result.model_dump(mode="json"))) - cur.execute( + async with asyncpg_pool.acquire() as conn: + # Store main enrichment result + results_escaped = json.dumps(helpers.sanitize_for_jsonb(result.model_dump(mode="json"))) + await conn.execute( + """ + INSERT INTO enrichments (object_id, module_name, result_data) + VALUES ($1, $2, $3) + """, + enrichment_request.object_id, + enrichment_name, + results_escaped, + ) + + # Store any transforms + if result.transforms: + for transform in result.transforms: + await conn.execute( """ - INSERT INTO enrichments (object_id, module_name, result_data) - VALUES (%s, %s, %s) + INSERT INTO transforms (object_id, type, transform_object_id, metadata) + VALUES ($1, $2, $3, $4) """, - (request.object_id, enrichment_name, results_escaped), + enrichment_request.object_id, + transform.type, + transform.object_id, + json.dumps(transform.metadata) if transform.metadata else None, ) - # Store any transforms - if result.transforms: - for transform in result.transforms: - cur.execute( - """ - INSERT INTO transforms (object_id, type, transform_object_id, metadata) - VALUES (%s, %s, %s, %s) - """, - ( - request.object_id, - transform.type, - transform.object_id, - json.dumps(transform.metadata) if transform.metadata else None, - ), - ) - - # Store any findings - if result.findings: - for finding in result.findings: - cur.execute( - """ - INSERT INTO findings ( - finding_name, category, severity, object_id, - origin_type, origin_name, raw_data, data - ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s) - """, - ( - finding.finding_name, - finding.category, - finding.severity, - request.object_id, - finding.origin_type, - finding.origin_name, - json.dumps(finding.raw_data), - json.dumps([obj.model_dump_json() for obj in finding.data]), - ), - ) - - conn.commit() - - # Run database operations in thread - await asyncio.to_thread(store_results) + # Store any findings + if result.findings: + for finding in result.findings: + await conn.execute( + """ + INSERT INTO findings ( + finding_name, category, severity, object_id, + origin_type, origin_name, raw_data, data + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8) + """, + finding.finding_name, + finding.category, + finding.severity, + enrichment_request.object_id, + finding.origin_type, + finding.origin_name, + json.dumps(finding.raw_data), + json.dumps([obj.model_dump_json() for obj in finding.data]), + ) return { "status": "success", "message": f"Completed enrichment with module '{enrichment_name}'", - "object_id": request.object_id, + "object_id": enrichment_request.object_id, "instance_id": str(uuid.uuid4()), # Generate a unique instance ID } @@ -162,7 +154,7 @@ async def run_enrichment( e, message="Error running enrichment module", enrichment_name=enrichment_name, - object_id=request.object_id, + object_id=enrichment_request.object_id, pid=os.getpid(), ) raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") from e diff --git a/projects/file_enrichment/file_enrichment/subscriptions/dotnet_handler.py b/projects/file_enrichment/file_enrichment/subscriptions/dotnet_handler.py index e7ebdb1..7768520 100644 --- a/projects/file_enrichment/file_enrichment/subscriptions/dotnet_handler.py +++ b/projects/file_enrichment/file_enrichment/subscriptions/dotnet_handler.py @@ -3,16 +3,16 @@ import json import os +import asyncpg from common.logger import get_logger from common.models import DotNetOutput from common.state_helpers import get_file_enriched_async from file_enrichment.dotnet import store_dotnet_results -from psycopg_pool import ConnectionPool logger = get_logger(__name__) -async def process_dotnet_event(raw_data, pool: ConnectionPool): +async def process_dotnet_event(raw_data, pool: asyncpg.Pool): """Process incoming .NET processing results from the dotnet_service""" try: logger.debug(f"Received DotNet output event: {raw_data}", pid=os.getpid()) @@ -36,10 +36,10 @@ async def process_dotnet_event(raw_data, pool: ConnectionPool): logger.debug(f"Processing dotnet results for object {object_id}", pid=os.getpid()) # Get the file enriched data for creating transforms - file_enriched = None try: file_enriched = await get_file_enriched_async(object_id) except Exception as e: + file_enriched = None logger.warning(f"Could not get file_enriched for {object_id}: {e}", pid=os.getpid()) # Store the results in the database using our helper function diff --git a/projects/file_enrichment/file_enrichment/subscriptions/file_handler.py b/projects/file_enrichment/file_enrichment/subscriptions/file_handler.py index e7ef01c..01adae2 100644 --- a/projects/file_enrichment/file_enrichment/subscriptions/file_handler.py +++ b/projects/file_enrichment/file_enrichment/subscriptions/file_handler.py @@ -1,17 +1,16 @@ """Handler for file subscription events.""" -import asyncio import os from datetime import datetime +import asyncpg from common.logger import get_logger from common.models import File -from psycopg_pool import ConnectionPool logger = get_logger(__name__) -async def save_file_message(file: File, pool: ConnectionPool): +async def save_file_message(file: File, pool: asyncpg.Pool): """Save the file message to the database for recovery purposes""" try: # Only save files that are not nested (originating files) @@ -24,53 +23,47 @@ async def save_file_message(file: File, pool: ConnectionPool): ) return - def save_to_db(): - with pool.connection() as conn: - with conn.cursor() as cur: - query = """ - INSERT INTO files ( - object_id, agent_id, source, project, timestamp, expiration, - path, originating_object_id, originating_container_id, nesting_level, - file_creation_time, file_access_time, file_modification_time - ) VALUES ( - %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s - ) ON CONFLICT (object_id) DO UPDATE SET - agent_id = EXCLUDED.agent_id, - source = EXCLUDED.source, - project = EXCLUDED.project, - timestamp = EXCLUDED.timestamp, - expiration = EXCLUDED.expiration, - path = EXCLUDED.path, - originating_object_id = EXCLUDED.originating_object_id, - originating_container_id = EXCLUDED.originating_container_id, - nesting_level = EXCLUDED.nesting_level, - file_creation_time = EXCLUDED.file_creation_time, - file_access_time = EXCLUDED.file_access_time, - file_modification_time = EXCLUDED.file_modification_time, - updated_at = CURRENT_TIMESTAMP; - """ + query = """ + INSERT INTO files ( + object_id, agent_id, source, project, timestamp, expiration, + path, originating_object_id, originating_container_id, nesting_level, + file_creation_time, file_access_time, file_modification_time + ) VALUES ( + $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13 + ) ON CONFLICT (object_id) DO UPDATE SET + agent_id = EXCLUDED.agent_id, + source = EXCLUDED.source, + project = EXCLUDED.project, + timestamp = EXCLUDED.timestamp, + expiration = EXCLUDED.expiration, + path = EXCLUDED.path, + originating_object_id = EXCLUDED.originating_object_id, + originating_container_id = EXCLUDED.originating_container_id, + nesting_level = EXCLUDED.nesting_level, + file_creation_time = EXCLUDED.file_creation_time, + file_access_time = EXCLUDED.file_access_time, + file_modification_time = EXCLUDED.file_modification_time, + updated_at = CURRENT_TIMESTAMP; + """ - cur.execute( - query, - ( - file.object_id, - file.agent_id, - file.source, - file.project, - file.timestamp, - file.expiration, - file.path, - file.originating_object_id, - getattr(file, "originating_container_id", None), - file.nesting_level, - datetime.fromisoformat(file.creation_time) if file.creation_time else None, - datetime.fromisoformat(file.access_time) if file.access_time else None, - datetime.fromisoformat(file.modification_time) if file.modification_time else None, - ), - ) - conn.commit() + async with pool.acquire() as conn: + await conn.execute( + query, + file.object_id, + file.agent_id, + file.source, + file.project, + file.timestamp, + file.expiration, + file.path, + file.originating_object_id, + getattr(file, "originating_container_id", None), + file.nesting_level, + datetime.fromisoformat(file.creation_time) if file.creation_time else None, + datetime.fromisoformat(file.access_time) if file.access_time else None, + datetime.fromisoformat(file.modification_time) if file.modification_time else None, + ) - await asyncio.to_thread(save_to_db) logger.debug("Successfully saved file message to database", object_id=file.object_id, pid=os.getpid()) except Exception as e: @@ -78,7 +71,7 @@ async def save_file_message(file: File, pool: ConnectionPool): raise -async def process_file_event(file: File, workflow_manager, module_execution_order: list, pool: ConnectionPool): +async def process_file_event(file: File, workflow_manager, module_execution_order: list, pool: asyncpg.Pool): """Process incoming file events""" try: # Save the file message to database first for recovery purposes diff --git a/projects/file_enrichment/file_enrichment/subscriptions/noseyparker_handler.py b/projects/file_enrichment/file_enrichment/subscriptions/noseyparker_handler.py index d2a3bcd..a29dad2 100644 --- a/projects/file_enrichment/file_enrichment/subscriptions/noseyparker_handler.py +++ b/projects/file_enrichment/file_enrichment/subscriptions/noseyparker_handler.py @@ -3,15 +3,15 @@ import json import os +import asyncpg from common.logger import get_logger from common.models import NoseyParkerOutput from file_enrichment.noseyparker import store_noseyparker_results -from psycopg_pool import ConnectionPool logger = get_logger(__name__) -async def process_noseyparker_event(raw_data, pool: ConnectionPool): +async def process_noseyparker_event(raw_data, pool: asyncpg.Pool): """Process incoming Nosey Parker scan results""" try: # Extract the raw data diff --git a/projects/file_enrichment/file_enrichment/workflow_recovery.py b/projects/file_enrichment/file_enrichment/workflow_recovery.py index d91e3f3..b5443ea 100644 --- a/projects/file_enrichment/file_enrichment/workflow_recovery.py +++ b/projects/file_enrichment/file_enrichment/workflow_recovery.py @@ -27,68 +27,18 @@ async def recover_interrupted_workflows(pool) -> None: logger.info("Starting workflow recovery process...", pid=os.getpid()) - def get_and_delete_running_workflows(): - with pool.connection() as conn: - with conn.cursor() as cur: - # Atomic DELETE with RETURNING - only one worker will get the interrupted workflows - cur.execute(""" - DELETE FROM workflows - WHERE status = 'RUNNING' - RETURNING object_id - """) - running_ids = [row[0] for row in cur.fetchall()] - conn.commit() + # Get interrupted workflows atomically using asyncpg + async with pool.acquire() as conn: + # Atomic DELETE with RETURNING - only one worker will get the interrupted workflows + running_ids = await conn.fetch(""" + DELETE FROM workflows + WHERE status = 'RUNNING' + RETURNING object_id + """) + running_object_ids = [row['object_id'] for row in running_ids] - if running_ids: - logger.info(f"Atomically claimed {len(running_ids)} interrupted workflows", pid=os.getpid()) - - return running_ids - - def get_file_data_and_cleanup(object_ids): - recovered_files = [] - with pool.connection() as conn: - with conn.cursor() as cur: - for object_id in object_ids: - # Get file data for reconstruction - cur.execute( - """ - SELECT object_id, agent_id, source, project, timestamp, expiration, - path, originating_object_id, originating_container_id, nesting_level, - file_creation_time, file_access_time, file_modification_time - FROM files WHERE object_id = %s - """, - (object_id,), - ) - - row = cur.fetchone() - if row: - # Convert database row to File-compatible dict - file_data = { - "object_id": str(row[0]), - "agent_id": row[1], - "source": row[2], - "project": row[3], - "timestamp": row[4], - "expiration": row[5], - "path": row[6], - "originating_object_id": str(row[7]) if row[7] else None, - "originating_container_id": str(row[8]) if row[8] else None, - "nesting_level": row[9], - "creation_time": row[10].isoformat() if row[10] else None, - "access_time": row[11].isoformat() if row[11] else None, - "modification_time": row[12].isoformat() if row[12] else None, - } - recovered_files.append(file_data) - logger.debug("Recovered file data for workflow", object_id=object_id, pid=os.getpid()) - else: - logger.warning("No file data found for workflow", object_id=object_id, pid=os.getpid()) - - conn.commit() - - return recovered_files - - # Get interrupted workflows - running_object_ids = await asyncio.to_thread(get_and_delete_running_workflows) + if running_object_ids: + logger.info(f"Atomically claimed {len(running_object_ids)} interrupted workflows", pid=os.getpid()) if not running_object_ids: logger.info("No interrupted workflows found", pid=os.getpid()) @@ -97,7 +47,41 @@ async def recover_interrupted_workflows(pool) -> None: logger.info(f"Found {len(running_object_ids)} interrupted workflows to recover", pid=os.getpid()) # Get file data and clean up partial results - recovered_files = await asyncio.to_thread(get_file_data_and_cleanup, running_object_ids) + recovered_files = [] + async with pool.acquire() as conn: + for object_id in running_object_ids: + # Get file data for reconstruction + row = await conn.fetchrow( + """ + SELECT object_id, agent_id, source, project, timestamp, expiration, + path, originating_object_id, originating_container_id, nesting_level, + file_creation_time, file_access_time, file_modification_time + FROM files WHERE object_id = $1 + """, + object_id, + ) + + if row: + # Convert database row to File-compatible dict + file_data = { + "object_id": str(row['object_id']), + "agent_id": row['agent_id'], + "source": row['source'], + "project": row['project'], + "timestamp": row['timestamp'], + "expiration": row['expiration'], + "path": row['path'], + "originating_object_id": str(row['originating_object_id']) if row['originating_object_id'] else None, + "originating_container_id": str(row['originating_container_id']) if row['originating_container_id'] else None, + "nesting_level": row['nesting_level'], + "creation_time": row['file_creation_time'].isoformat() if row['file_creation_time'] else None, + "access_time": row['file_access_time'].isoformat() if row['file_access_time'] else None, + "modification_time": row['file_modification_time'].isoformat() if row['file_modification_time'] else None, + } + recovered_files.append(file_data) + logger.debug("Recovered file data for workflow", object_id=object_id, pid=os.getpid()) + else: + logger.warning("No file data found for workflow", object_id=object_id, pid=os.getpid()) if not recovered_files: logger.warning("No file data found for interrupted workflows", pid=os.getpid())