Fix non-atomic async operation creation (#619)
* Fix non-atomic async operation creation in _submit_async_operation Previously the method performed two separate database round-trips: 1. INSERT into async_operations with no task_payload (null) 2. submit_task → UPDATE to set task_payload A process crash or network error between steps 1 and 2 left a row with task_payload IS NULL permanently. The worker's claim query requires task_payload IS NOT NULL, so these orphaned rows could never be picked up and the queue appeared degraded indefinitely. Fix: build full_payload before the INSERT and include task_payload in the same INSERT statement, making operation creation atomic. submit_task is still called afterwards — for SyncTaskBackend it executes the task immediately (unchanged behaviour); for BrokerTaskBackend it becomes an idempotent UPDATE (payload already set) kept for symmetry. * Preserve datetime payloads in atomic async insert
This commit is contained in:
parent
439424559e
commit
94cf89b570
2 changed files with 101 additions and 15 deletions
|
|
@ -67,6 +67,13 @@ def fq_table(table_name: str) -> str:
|
||||||
return f"{get_current_schema()}.{table_name}"
|
return f"{get_current_schema()}.{table_name}"
|
||||||
|
|
||||||
|
|
||||||
|
def _json_default(obj: Any) -> str:
|
||||||
|
"""JSON serializer for types commonly carried through async task payloads."""
|
||||||
|
if isinstance(obj, datetime):
|
||||||
|
return obj.isoformat()
|
||||||
|
raise TypeError(f"Object of type {type(obj).__name__} is not JSON serializable")
|
||||||
|
|
||||||
|
|
||||||
# Tables that must be schema-qualified (for runtime validation)
|
# Tables that must be schema-qualified (for runtime validation)
|
||||||
_PROTECTED_TABLES = frozenset(
|
_PROTECTED_TABLES = frozenset(
|
||||||
[
|
[
|
||||||
|
|
@ -7407,21 +7414,10 @@ class MemoryEngine(MemoryEngineInterface):
|
||||||
|
|
||||||
operation_id = uuid.uuid4()
|
operation_id = uuid.uuid4()
|
||||||
|
|
||||||
# Insert operation record into database
|
# Build full payload before INSERT so task_payload is included atomically.
|
||||||
async with acquire_with_retry(pool) as conn:
|
# Previously the INSERT omitted task_payload and a separate submit_task call
|
||||||
await conn.execute(
|
# did an UPDATE — a crash between the two left a null-payload row that the
|
||||||
f"""
|
# worker's claim query (task_payload IS NOT NULL) could never pick up.
|
||||||
INSERT INTO {fq_table("async_operations")} (operation_id, bank_id, operation_type, result_metadata, status)
|
|
||||||
VALUES ($1, $2, $3, $4, $5)
|
|
||||||
""",
|
|
||||||
operation_id,
|
|
||||||
bank_id,
|
|
||||||
operation_type,
|
|
||||||
json.dumps(result_metadata or {}),
|
|
||||||
"pending",
|
|
||||||
)
|
|
||||||
|
|
||||||
# Build and submit task payload
|
|
||||||
full_payload = {
|
full_payload = {
|
||||||
"type": task_type,
|
"type": task_type,
|
||||||
"operation_id": str(operation_id),
|
"operation_id": str(operation_id),
|
||||||
|
|
@ -7429,6 +7425,24 @@ class MemoryEngine(MemoryEngineInterface):
|
||||||
**task_payload,
|
**task_payload,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
# Insert operation record with task_payload in a single atomic statement
|
||||||
|
async with acquire_with_retry(pool) as conn:
|
||||||
|
await conn.execute(
|
||||||
|
f"""
|
||||||
|
INSERT INTO {fq_table("async_operations")} (operation_id, bank_id, operation_type, result_metadata, status, task_payload)
|
||||||
|
VALUES ($1, $2, $3, $4, $5, $6::jsonb)
|
||||||
|
""",
|
||||||
|
operation_id,
|
||||||
|
bank_id,
|
||||||
|
operation_type,
|
||||||
|
json.dumps(result_metadata or {}, default=_json_default),
|
||||||
|
"pending",
|
||||||
|
json.dumps(full_payload, default=_json_default),
|
||||||
|
)
|
||||||
|
|
||||||
|
# For SyncTaskBackend: executes the task immediately.
|
||||||
|
# For BrokerTaskBackend: does an idempotent UPDATE (payload already set above),
|
||||||
|
# kept for symmetry and to support any future notification mechanisms.
|
||||||
await self._task_backend.submit_task(full_payload)
|
await self._task_backend.submit_task(full_payload)
|
||||||
|
|
||||||
logger.info(f"{operation_type} task queued for bank_id={bank_id}, operation_id={operation_id}")
|
logger.info(f"{operation_type} task queued for bank_id={bank_id}, operation_id={operation_id}")
|
||||||
|
|
|
||||||
|
|
@ -5,6 +5,7 @@ End-to-end tests for file retain (upload, convert, retain) functionality.
|
||||||
import asyncio
|
import asyncio
|
||||||
import io
|
import io
|
||||||
import json
|
import json
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from httpx import ASGITransport, AsyncClient
|
from httpx import ASGITransport, AsyncClient
|
||||||
|
|
@ -471,6 +472,77 @@ async def test_file_conversion_creates_separate_retain_operation(memory_no_llm_v
|
||||||
assert len(doc["original_text"]) > 0
|
assert len(doc["original_text"]) > 0
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_file_retain_serializes_datetime_timestamp(memory_no_llm_verify, sample_txt_content):
|
||||||
|
"""Async file retain should accept Python datetimes in task payloads."""
|
||||||
|
from hindsight_api.engine.parsers.base import FileParser
|
||||||
|
from hindsight_api.models import RequestContext
|
||||||
|
|
||||||
|
bank_id = f"test_file_timestamp_bank_{datetime.now(timezone.utc).timestamp()}"
|
||||||
|
timestamp = datetime(2024, 1, 15, 10, 30, tzinfo=timezone.utc)
|
||||||
|
|
||||||
|
context = RequestContext(internal=True)
|
||||||
|
await memory_no_llm_verify.get_bank_profile(bank_id, request_context=context)
|
||||||
|
|
||||||
|
class MockFile:
|
||||||
|
def __init__(self, content, filename, content_type):
|
||||||
|
self.content = content
|
||||||
|
self.filename = filename
|
||||||
|
self.content_type = content_type
|
||||||
|
|
||||||
|
async def read(self):
|
||||||
|
return self.content
|
||||||
|
|
||||||
|
class TimestampParser(FileParser):
|
||||||
|
async def convert(self, file_data: bytes, filename: str) -> str:
|
||||||
|
return file_data.decode("utf-8")
|
||||||
|
|
||||||
|
def supports(self, filename: str, content_type: str | None = None) -> bool:
|
||||||
|
return filename.endswith(".txt")
|
||||||
|
|
||||||
|
def name(self) -> str:
|
||||||
|
return "timestamp_parser"
|
||||||
|
|
||||||
|
memory_no_llm_verify._parser_registry.register(TimestampParser())
|
||||||
|
|
||||||
|
mock_file = MockFile(sample_txt_content, "timestamped.txt", "text/plain")
|
||||||
|
|
||||||
|
result = await memory_no_llm_verify.submit_async_file_retain(
|
||||||
|
bank_id=bank_id,
|
||||||
|
file_items=[
|
||||||
|
{
|
||||||
|
"file": mock_file,
|
||||||
|
"document_id": "timestamped_doc",
|
||||||
|
"context": "timestamp test",
|
||||||
|
"metadata": {},
|
||||||
|
"tags": [],
|
||||||
|
"timestamp": timestamp,
|
||||||
|
"parser": ["timestamp_parser"],
|
||||||
|
}
|
||||||
|
],
|
||||||
|
document_tags=None,
|
||||||
|
request_context=context,
|
||||||
|
)
|
||||||
|
|
||||||
|
operation_id = result["operation_ids"][0]
|
||||||
|
pool = await memory_no_llm_verify._get_pool()
|
||||||
|
from hindsight_api.engine.memory_engine import get_current_schema
|
||||||
|
|
||||||
|
async with pool.acquire() as conn:
|
||||||
|
row = await conn.fetchrow(
|
||||||
|
f"""
|
||||||
|
SELECT status, task_payload->>'timestamp' AS timestamp
|
||||||
|
FROM {get_current_schema()}.async_operations
|
||||||
|
WHERE operation_id = $1
|
||||||
|
""",
|
||||||
|
operation_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert row is not None
|
||||||
|
assert row["status"] == "completed"
|
||||||
|
assert row["timestamp"] == "2024-01-15T10:30:00+00:00"
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_file_conversion_failure_sets_status_to_failed(memory_no_llm_verify, sample_txt_content):
|
async def test_file_conversion_failure_sets_status_to_failed(memory_no_llm_verify, sample_txt_content):
|
||||||
"""Test that when file conversion fails, the operation status is set to 'failed' not 'completed'."""
|
"""Test that when file conversion fails, the operation status is set to 'failed' not 'completed'."""
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue