3fc87e767c
compute_semantic_links_ann created a TEMP TABLE outside any transaction, then ran a TRUNCATE / COPY / SELECT / DROP sequence as separate statements on the same asyncpg connection. This is fine against a direct Postgres connection but fails intermittently when the caller is routed through PgBouncer in transaction pool mode: CREATE TEMP TABLE IF NOT EXISTS _ann_seeds (...) -- backend A TRUNCATE _ann_seeds -- backend B -> FAILS Temp tables are session-scoped to the backend that created them. In PgBouncer transaction mode the backend is only pinned to the client for the duration of an actual transaction, so between standalone statements the pooler can (and under concurrency, will) rebind the client to a different backend. When that happens the _ann_seeds table disappears and the follow-up statement fails with: relation "_ann_seeds" does not exist Symptom: ~3% of sync retain calls (2 of 61) failed the Hindsight Cloud smoke test on a recent hindsight-dev deploy. Async retains are masked by the 3-attempt retry loop so they usually eventually succeed. Fix: wrap the CREATE TEMP TABLE -> COPY -> SELECT sequence in a single `async with conn.transaction():` block, and use ON COMMIT DROP so the temp table is transaction-scoped and auto-cleaned at commit. Also switch `SET hnsw.ef_search = 60` to `SET LOCAL` so the tuning is transaction-scoped and no longer leaks onto the pooled backend for subsequent recall queries. Drop the now-unnecessary manual TRUNCATE, explicit DROP TABLE, and RESET hnsw.ef_search. The function docstring still correctly describes this as running on a separate connection outside the surrounding write transaction — this change only adds an inner transaction around the ANN work itself to keep the temp table visible to PgBouncer. Tests: - Add TestComputeSemanticLinksAnnPgBouncerSafety with 5 regression tests using a mocked connection. These are structural asserts — they check that the function enters conn.transaction(), uses ON COMMIT DROP, uses SET LOCAL, and does not reintroduce manual TRUNCATE / DROP / RESET calls. They would have caught the original bug if they had existed, and will catch any future reversion.
534 lines
22 KiB
Python
534 lines
22 KiB
Python
"""Tests for link_utils datetime handling, temporal link computation, and semantic link splitting."""
|
|
import numpy as np
|
|
import pytest
|
|
from datetime import datetime, timezone, timedelta
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
from hindsight_api.engine.retain.link_utils import (
|
|
_normalize_datetime,
|
|
_cap_links_per_unit,
|
|
compute_temporal_links,
|
|
compute_temporal_query_bounds,
|
|
compute_semantic_links_ann,
|
|
compute_semantic_links_within_batch,
|
|
MAX_TEMPORAL_LINKS_PER_UNIT,
|
|
)
|
|
|
|
|
|
class TestNormalizeDatetime:
|
|
"""Tests for the _normalize_datetime helper function."""
|
|
|
|
def test_none_returns_none(self):
|
|
"""Test that None input returns None."""
|
|
assert _normalize_datetime(None) is None
|
|
|
|
def test_naive_datetime_becomes_utc(self):
|
|
"""Test that naive datetimes are converted to UTC."""
|
|
naive_dt = datetime(2024, 6, 15, 10, 30, 0)
|
|
result = _normalize_datetime(naive_dt)
|
|
|
|
assert result.tzinfo is not None
|
|
assert result.tzinfo == timezone.utc
|
|
assert result.year == 2024
|
|
assert result.month == 6
|
|
assert result.day == 15
|
|
assert result.hour == 10
|
|
assert result.minute == 30
|
|
|
|
def test_aware_datetime_unchanged(self):
|
|
"""Test that timezone-aware datetimes are returned unchanged."""
|
|
aware_dt = datetime(2024, 6, 15, 10, 30, 0, tzinfo=timezone.utc)
|
|
result = _normalize_datetime(aware_dt)
|
|
|
|
assert result == aware_dt
|
|
assert result.tzinfo == timezone.utc
|
|
|
|
def test_mixed_datetimes_can_be_compared(self):
|
|
"""Test that normalized naive and aware datetimes can be compared."""
|
|
naive_dt = datetime(2024, 6, 15, 10, 30, 0)
|
|
aware_dt = datetime(2024, 6, 15, 10, 30, 0, tzinfo=timezone.utc)
|
|
|
|
normalized_naive = _normalize_datetime(naive_dt)
|
|
normalized_aware = _normalize_datetime(aware_dt)
|
|
|
|
# Should be able to compare without TypeError
|
|
assert normalized_naive == normalized_aware
|
|
|
|
|
|
class TestComputeTemporalQueryBounds:
|
|
"""Tests for compute_temporal_query_bounds function."""
|
|
|
|
def test_empty_units_returns_none(self):
|
|
"""Test that empty input returns (None, None)."""
|
|
min_date, max_date = compute_temporal_query_bounds({})
|
|
assert min_date is None
|
|
assert max_date is None
|
|
|
|
def test_single_unit_normal_date(self):
|
|
"""Test bounds for a single unit with normal date."""
|
|
units = {"unit-1": datetime(2024, 6, 15, 12, 0, 0, tzinfo=timezone.utc)}
|
|
min_date, max_date = compute_temporal_query_bounds(units, time_window_hours=24)
|
|
|
|
assert min_date == datetime(2024, 6, 14, 12, 0, 0, tzinfo=timezone.utc)
|
|
assert max_date == datetime(2024, 6, 16, 12, 0, 0, tzinfo=timezone.utc)
|
|
|
|
def test_multiple_units(self):
|
|
"""Test bounds span across multiple units."""
|
|
units = {
|
|
"unit-1": datetime(2024, 6, 10, 12, 0, 0, tzinfo=timezone.utc),
|
|
"unit-2": datetime(2024, 6, 15, 12, 0, 0, tzinfo=timezone.utc),
|
|
"unit-3": datetime(2024, 6, 20, 12, 0, 0, tzinfo=timezone.utc),
|
|
}
|
|
min_date, max_date = compute_temporal_query_bounds(units, time_window_hours=24)
|
|
|
|
# min should be Jun 10 - 24h = Jun 9
|
|
assert min_date == datetime(2024, 6, 9, 12, 0, 0, tzinfo=timezone.utc)
|
|
# max should be Jun 20 + 24h = Jun 21
|
|
assert max_date == datetime(2024, 6, 21, 12, 0, 0, tzinfo=timezone.utc)
|
|
|
|
def test_mixed_naive_and_aware_datetimes(self):
|
|
"""Test that mixed naive/aware datetimes work correctly."""
|
|
units = {
|
|
"unit-1": datetime(2024, 6, 10, 12, 0, 0), # naive
|
|
"unit-2": datetime(2024, 6, 15, 12, 0, 0, tzinfo=timezone.utc), # aware
|
|
}
|
|
# Should not raise TypeError
|
|
min_date, max_date = compute_temporal_query_bounds(units, time_window_hours=24)
|
|
|
|
assert min_date is not None
|
|
assert max_date is not None
|
|
assert min_date.tzinfo is not None
|
|
assert max_date.tzinfo is not None
|
|
|
|
def test_overflow_near_datetime_min(self):
|
|
"""Test overflow protection near datetime.min."""
|
|
units = {"unit-1": datetime(1, 1, 2, 0, 0, tzinfo=timezone.utc)}
|
|
min_date, max_date = compute_temporal_query_bounds(units, time_window_hours=48)
|
|
|
|
# Should handle overflow gracefully
|
|
assert min_date == datetime.min.replace(tzinfo=timezone.utc)
|
|
assert max_date is not None
|
|
|
|
def test_overflow_near_datetime_max(self):
|
|
"""Test overflow protection near datetime.max."""
|
|
units = {"unit-1": datetime(9999, 12, 30, 0, 0, tzinfo=timezone.utc)}
|
|
min_date, max_date = compute_temporal_query_bounds(units, time_window_hours=48)
|
|
|
|
# Should handle overflow gracefully
|
|
assert min_date is not None
|
|
assert max_date == datetime.max.replace(tzinfo=timezone.utc)
|
|
|
|
|
|
class TestComputeTemporalLinks:
|
|
"""Tests for compute_temporal_links function."""
|
|
|
|
def test_empty_units_returns_empty(self):
|
|
"""Test that empty input returns empty list."""
|
|
links = compute_temporal_links({}, [])
|
|
assert links == []
|
|
|
|
def test_no_candidates_returns_empty(self):
|
|
"""Test that no candidates means no links."""
|
|
units = {"unit-1": datetime(2024, 6, 15, 12, 0, 0, tzinfo=timezone.utc)}
|
|
links = compute_temporal_links(units, [])
|
|
assert links == []
|
|
|
|
def test_candidate_within_window_creates_link(self):
|
|
"""Test that candidates within time window create links."""
|
|
units = {"unit-1": datetime(2024, 6, 15, 12, 0, 0, tzinfo=timezone.utc)}
|
|
candidates = [
|
|
{"id": "candidate-1", "event_date": datetime(2024, 6, 15, 10, 0, 0, tzinfo=timezone.utc)},
|
|
]
|
|
|
|
links = compute_temporal_links(units, candidates, time_window_hours=24)
|
|
|
|
assert len(links) == 1
|
|
assert links[0][0] == "unit-1"
|
|
assert links[0][1] == "candidate-1"
|
|
assert links[0][2] == "temporal"
|
|
assert links[0][4] is None
|
|
# Weight should be high since they're close (2 hours apart)
|
|
assert links[0][3] > 0.9
|
|
|
|
def test_candidate_outside_window_no_link(self):
|
|
"""Test that candidates outside time window don't create links."""
|
|
units = {"unit-1": datetime(2024, 6, 15, 12, 0, 0, tzinfo=timezone.utc)}
|
|
candidates = [
|
|
{"id": "candidate-1", "event_date": datetime(2024, 6, 10, 12, 0, 0, tzinfo=timezone.utc)},
|
|
]
|
|
|
|
links = compute_temporal_links(units, candidates, time_window_hours=24)
|
|
|
|
assert len(links) == 0
|
|
|
|
def test_weight_decreases_with_distance(self):
|
|
"""Test that weight decreases as time difference increases."""
|
|
units = {"unit-1": datetime(2024, 6, 15, 12, 0, 0, tzinfo=timezone.utc)}
|
|
candidates = [
|
|
{"id": "close", "event_date": datetime(2024, 6, 15, 11, 0, 0, tzinfo=timezone.utc)}, # 1 hour
|
|
{"id": "far", "event_date": datetime(2024, 6, 14, 18, 0, 0, tzinfo=timezone.utc)}, # 18 hours
|
|
]
|
|
|
|
links = compute_temporal_links(units, candidates, time_window_hours=24)
|
|
|
|
assert len(links) == 2
|
|
close_link = next(l for l in links if l[1] == "close")
|
|
far_link = next(l for l in links if l[1] == "far")
|
|
|
|
assert close_link[3] > far_link[3]
|
|
|
|
def test_max_10_links_per_unit(self):
|
|
"""Test that at most 10 links are created per unit."""
|
|
units = {"unit-1": datetime(2024, 6, 15, 12, 0, 0, tzinfo=timezone.utc)}
|
|
# Create 15 candidates all within window
|
|
candidates = [
|
|
{"id": f"candidate-{i}", "event_date": datetime(2024, 6, 15, 11, 0, 0, tzinfo=timezone.utc)}
|
|
for i in range(15)
|
|
]
|
|
|
|
links = compute_temporal_links(units, candidates, time_window_hours=24)
|
|
|
|
assert len(links) == 10
|
|
|
|
def test_multiple_units_multiple_candidates(self):
|
|
"""Test with multiple units and candidates."""
|
|
units = {
|
|
"unit-1": datetime(2024, 6, 15, 12, 0, 0, tzinfo=timezone.utc),
|
|
"unit-2": datetime(2024, 6, 20, 12, 0, 0, tzinfo=timezone.utc),
|
|
}
|
|
candidates = [
|
|
{"id": "c1", "event_date": datetime(2024, 6, 15, 10, 0, 0, tzinfo=timezone.utc)}, # near unit-1
|
|
{"id": "c2", "event_date": datetime(2024, 6, 20, 10, 0, 0, tzinfo=timezone.utc)}, # near unit-2
|
|
{"id": "c3", "event_date": datetime(2024, 6, 17, 12, 0, 0, tzinfo=timezone.utc)}, # between, near neither
|
|
]
|
|
|
|
links = compute_temporal_links(units, candidates, time_window_hours=24)
|
|
|
|
# unit-1 should link to c1 only
|
|
# unit-2 should link to c2 only
|
|
unit1_links = [l for l in links if l[0] == "unit-1"]
|
|
unit2_links = [l for l in links if l[0] == "unit-2"]
|
|
|
|
assert len(unit1_links) == 1
|
|
assert unit1_links[0][1] == "c1"
|
|
|
|
assert len(unit2_links) == 1
|
|
assert unit2_links[0][1] == "c2"
|
|
|
|
def test_mixed_naive_and_aware_datetimes(self):
|
|
"""Test that mixed naive/aware datetimes work correctly."""
|
|
units = {"unit-1": datetime(2024, 6, 15, 12, 0, 0)} # naive
|
|
candidates = [
|
|
{"id": "c1", "event_date": datetime(2024, 6, 15, 10, 0, 0, tzinfo=timezone.utc)}, # aware
|
|
]
|
|
|
|
# Should not raise TypeError
|
|
links = compute_temporal_links(units, candidates, time_window_hours=24)
|
|
assert len(links) == 1
|
|
|
|
def test_overflow_near_datetime_min(self):
|
|
"""Test overflow protection when unit date is near datetime.min."""
|
|
units = {"unit-1": datetime(1, 1, 2, 0, 0, tzinfo=timezone.utc)}
|
|
candidates = [
|
|
{"id": "c1", "event_date": datetime(1, 1, 1, 12, 0, 0, tzinfo=timezone.utc)},
|
|
]
|
|
|
|
# Should not raise OverflowError
|
|
links = compute_temporal_links(units, candidates, time_window_hours=48)
|
|
assert len(links) == 1
|
|
|
|
def test_overflow_near_datetime_max(self):
|
|
"""Test overflow protection when unit date is near datetime.max."""
|
|
units = {"unit-1": datetime(9999, 12, 30, 0, 0, tzinfo=timezone.utc)}
|
|
candidates = [
|
|
{"id": "c1", "event_date": datetime(9999, 12, 31, 12, 0, 0, tzinfo=timezone.utc)},
|
|
]
|
|
|
|
# Should not raise OverflowError
|
|
links = compute_temporal_links(units, candidates, time_window_hours=48)
|
|
assert len(links) == 1
|
|
|
|
def test_weight_minimum_is_0_3(self):
|
|
"""Test that weight doesn't go below 0.3."""
|
|
units = {"unit-1": datetime(2024, 6, 15, 12, 0, 0, tzinfo=timezone.utc)}
|
|
candidates = [
|
|
# 23 hours apart - should be just within 24h window but low weight
|
|
{"id": "c1", "event_date": datetime(2024, 6, 14, 13, 0, 0, tzinfo=timezone.utc)},
|
|
]
|
|
|
|
links = compute_temporal_links(units, candidates, time_window_hours=24)
|
|
|
|
assert len(links) == 1
|
|
assert links[0][3] >= 0.3
|
|
|
|
|
|
class TestCapLinksPerUnit:
|
|
"""Tests for the _cap_links_per_unit helper function."""
|
|
|
|
def test_empty_links(self):
|
|
assert _cap_links_per_unit([]) == []
|
|
|
|
def test_under_cap_unchanged(self):
|
|
links = [
|
|
("unit_a", "unit_x", "temporal", 0.9, None),
|
|
("unit_a", "unit_y", "temporal", 0.8, None),
|
|
]
|
|
result = _cap_links_per_unit(links, max_per_unit=5)
|
|
assert len(result) == 2
|
|
|
|
def test_caps_to_max_per_unit(self):
|
|
# Create 30 links from the same unit with descending weights
|
|
links = [("unit_a", f"unit_{i}", "temporal", 1.0 - i * 0.01, None) for i in range(30)]
|
|
result = _cap_links_per_unit(links, max_per_unit=10)
|
|
assert len(result) == 10
|
|
# Should keep the highest-weight links
|
|
weights = [lnk[3] for lnk in result]
|
|
assert weights == sorted(weights, reverse=True)
|
|
assert weights[0] == 1.0 # Highest weight kept
|
|
|
|
def test_caps_independently_per_unit(self):
|
|
links_a = [("unit_a", f"target_{i}", "temporal", 0.9 - i * 0.01, None) for i in range(10)]
|
|
links_b = [("unit_b", f"target_{i}", "temporal", 0.8 - i * 0.01, None) for i in range(10)]
|
|
result = _cap_links_per_unit(links_a + links_b, max_per_unit=5)
|
|
# 5 from unit_a + 5 from unit_b
|
|
assert len(result) == 10
|
|
from_a = [lnk for lnk in result if lnk[0] == "unit_a"]
|
|
from_b = [lnk for lnk in result if lnk[0] == "unit_b"]
|
|
assert len(from_a) == 5
|
|
assert len(from_b) == 5
|
|
|
|
def test_default_max_is_temporal_constant(self):
|
|
links = [("unit_a", f"target_{i}", "temporal", 1.0 - i * 0.01, None) for i in range(50)]
|
|
result = _cap_links_per_unit(links)
|
|
assert len(result) == MAX_TEMPORAL_LINKS_PER_UNIT
|
|
|
|
def test_preserves_tuple_structure(self):
|
|
links = [("from_id", "to_id", "temporal", 0.95, "entity_id")]
|
|
result = _cap_links_per_unit(links, max_per_unit=5)
|
|
assert result[0] == ("from_id", "to_id", "temporal", 0.95, "entity_id")
|
|
|
|
|
|
class TestComputeSemanticLinksWithinBatch:
|
|
"""Tests for compute_semantic_links_within_batch.
|
|
|
|
This function computes semantic links between units in the same batch
|
|
using numpy dot product (no DB access). It runs in Phase 2 (write
|
|
transaction) while the expensive ANN search against existing units runs
|
|
in Phase 1 on a separate connection to avoid TimeoutErrors from HNSW
|
|
index contention under concurrent load.
|
|
"""
|
|
|
|
def test_empty_returns_empty(self):
|
|
assert compute_semantic_links_within_batch([], []) == []
|
|
|
|
def test_single_unit_returns_empty(self):
|
|
emb = [np.random.randn(384).tolist()]
|
|
assert compute_semantic_links_within_batch(["u1"], emb) == []
|
|
|
|
def test_identical_embeddings_produce_links(self):
|
|
"""Two identical embeddings should have similarity=1.0 (above 0.7 threshold)."""
|
|
emb = [0.1] * 384
|
|
links = compute_semantic_links_within_batch(["u1", "u2"], [emb, emb])
|
|
assert len(links) == 2 # bidirectional: u1→u2, u2→u1
|
|
from_ids = {lnk[0] for lnk in links}
|
|
to_ids = {lnk[1] for lnk in links}
|
|
assert from_ids == {"u1", "u2"}
|
|
assert to_ids == {"u1", "u2"}
|
|
for lnk in links:
|
|
assert lnk[2] == "semantic"
|
|
assert lnk[3] >= 0.99 # near-1.0 similarity
|
|
assert lnk[4] is None # no entity_id
|
|
|
|
def test_orthogonal_embeddings_no_links(self):
|
|
"""Orthogonal embeddings should have similarity=0 (below 0.7 threshold)."""
|
|
emb1 = [1.0] + [0.0] * 383
|
|
emb2 = [0.0] + [1.0] + [0.0] * 382
|
|
links = compute_semantic_links_within_batch(["u1", "u2"], [emb1, emb2])
|
|
assert len(links) == 0
|
|
|
|
def test_respects_threshold(self):
|
|
"""Links below threshold should be excluded."""
|
|
emb1 = np.random.randn(384).tolist()
|
|
# Create a slightly similar embedding (add noise)
|
|
emb2 = [x + np.random.randn() * 0.5 for x in emb1]
|
|
# Normalize both
|
|
norm1 = np.linalg.norm(emb1)
|
|
norm2 = np.linalg.norm(emb2)
|
|
emb1 = [x / norm1 for x in emb1]
|
|
emb2 = [x / norm2 for x in emb2]
|
|
|
|
links_low = compute_semantic_links_within_batch(["u1", "u2"], [emb1, emb2], threshold=0.0)
|
|
links_high = compute_semantic_links_within_batch(["u1", "u2"], [emb1, emb2], threshold=0.99)
|
|
# Low threshold should have more links than high threshold
|
|
assert len(links_low) >= len(links_high)
|
|
|
|
def test_top_k_limits_per_unit(self):
|
|
"""Each unit should link to at most top_k other units."""
|
|
n = 10
|
|
# Create similar embeddings (all close to the same vector)
|
|
base = np.random.randn(384)
|
|
base = base / np.linalg.norm(base)
|
|
embs = [(base + np.random.randn(384) * 0.01).tolist() for _ in range(n)]
|
|
unit_ids = [f"u{i}" for i in range(n)]
|
|
|
|
links = compute_semantic_links_within_batch(unit_ids, embs, top_k=3, threshold=0.5)
|
|
# Each unit should have at most 3 outgoing links
|
|
from collections import Counter
|
|
from_counts = Counter(lnk[0] for lnk in links)
|
|
for count in from_counts.values():
|
|
assert count <= 3
|
|
|
|
def test_link_tuple_structure(self):
|
|
"""Verify the tuple format matches what _bulk_insert_links expects."""
|
|
emb = [0.1] * 384
|
|
links = compute_semantic_links_within_batch(["u1", "u2"], [emb, emb])
|
|
for lnk in links:
|
|
assert len(lnk) == 5
|
|
from_id, to_id, link_type, weight, entity_id = lnk
|
|
assert isinstance(from_id, str)
|
|
assert isinstance(to_id, str)
|
|
assert link_type == "semantic"
|
|
assert 0.0 <= weight <= 1.0
|
|
assert entity_id is None
|
|
|
|
|
|
class TestComputeSemanticLinksAnnPgBouncerSafety:
|
|
"""Regression tests ensuring compute_semantic_links_ann stays in a single
|
|
transaction so that the `_ann_seeds` temp table remains visible when the
|
|
caller's connection goes through pgBouncer in `transaction` pool mode.
|
|
|
|
In pgBouncer transaction mode, the backend is only pinned to the client
|
|
for the duration of an actual PostgreSQL transaction. Outside a
|
|
transaction, consecutive statements can land on different backends, and
|
|
session-scoped temp tables (which are bound to the backend that created
|
|
them) become invisible. The observed failure mode was an intermittent
|
|
`relation "_ann_seeds" does not exist` on the statement immediately
|
|
following the CREATE TEMP TABLE.
|
|
"""
|
|
|
|
@pytest.fixture
|
|
def mock_conn(self):
|
|
"""An asyncpg-like connection mock with an async `transaction()`
|
|
context manager and awaitable execute/fetch/copy helpers."""
|
|
conn = MagicMock()
|
|
|
|
txn_cm = MagicMock()
|
|
txn_cm.__aenter__ = AsyncMock(return_value=None)
|
|
txn_cm.__aexit__ = AsyncMock(return_value=None)
|
|
conn.transaction = MagicMock(return_value=txn_cm)
|
|
|
|
conn.execute = AsyncMock()
|
|
conn.copy_records_to_table = AsyncMock()
|
|
conn.fetch = AsyncMock(return_value=[])
|
|
return conn
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_empty_inputs_skip_transaction(self, mock_conn):
|
|
"""No seeds -> no work, no transaction, no temp-table churn."""
|
|
result = await compute_semantic_links_ann(
|
|
conn=mock_conn,
|
|
bank_id="bank-1",
|
|
unit_ids=[],
|
|
embeddings=[],
|
|
)
|
|
assert result == []
|
|
mock_conn.transaction.assert_not_called()
|
|
mock_conn.execute.assert_not_called()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_runs_inside_a_transaction(self, mock_conn):
|
|
"""The full CREATE TEMP TABLE -> COPY -> SELECT sequence must happen
|
|
inside a single `async with conn.transaction():` block."""
|
|
emb = [0.1] * 384
|
|
await compute_semantic_links_ann(
|
|
conn=mock_conn,
|
|
bank_id="bank-1",
|
|
unit_ids=["u1", "u2"],
|
|
embeddings=[emb, emb],
|
|
fact_types=["world", "world"],
|
|
)
|
|
|
|
# Transaction context manager was entered.
|
|
mock_conn.transaction.assert_called_once()
|
|
txn_cm = mock_conn.transaction.return_value
|
|
txn_cm.__aenter__.assert_awaited_once()
|
|
txn_cm.__aexit__.assert_awaited_once()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_temp_table_uses_on_commit_drop(self, mock_conn):
|
|
"""The CREATE TEMP TABLE statement must use ON COMMIT DROP so the
|
|
table is transaction-scoped. Without ON COMMIT DROP the table would
|
|
be session-scoped and would not survive pgBouncer backend rebinding
|
|
between transactions."""
|
|
emb = [0.1] * 384
|
|
await compute_semantic_links_ann(
|
|
conn=mock_conn,
|
|
bank_id="bank-1",
|
|
unit_ids=["u1"],
|
|
embeddings=[emb],
|
|
fact_types=["world"],
|
|
)
|
|
|
|
executed_sql = [call.args[0] for call in mock_conn.execute.call_args_list]
|
|
create_statements = [s for s in executed_sql if "CREATE TEMP TABLE" in s]
|
|
assert len(create_statements) == 1, "Should create _ann_seeds exactly once"
|
|
assert "_ann_seeds" in create_statements[0]
|
|
assert "ON COMMIT DROP" in create_statements[0], (
|
|
"CREATE TEMP TABLE must use ON COMMIT DROP so the table is cleaned "
|
|
"up at transaction end and is transaction-scoped"
|
|
)
|
|
|
|
# Must not use IF NOT EXISTS — the table is fresh each transaction.
|
|
assert "IF NOT EXISTS" not in create_statements[0], (
|
|
"With ON COMMIT DROP the table is always fresh at transaction start, "
|
|
"so IF NOT EXISTS is both unnecessary and misleading (suggests the "
|
|
"table might persist across transactions)"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_manual_drop_or_truncate(self, mock_conn):
|
|
"""With ON COMMIT DROP we must not re-add manual TRUNCATE or DROP
|
|
statements — they were the source of the original pgBouncer bug."""
|
|
emb = [0.1] * 384
|
|
await compute_semantic_links_ann(
|
|
conn=mock_conn,
|
|
bank_id="bank-1",
|
|
unit_ids=["u1"],
|
|
embeddings=[emb],
|
|
fact_types=["world"],
|
|
)
|
|
|
|
executed_sql = [call.args[0] for call in mock_conn.execute.call_args_list]
|
|
assert not any("TRUNCATE _ann_seeds" in s for s in executed_sql), (
|
|
"TRUNCATE is unnecessary with ON COMMIT DROP and was previously "
|
|
"the statement that failed with 'relation does not exist' when "
|
|
"pgBouncer rebound the backend"
|
|
)
|
|
assert not any("DROP TABLE" in s and "_ann_seeds" in s for s in executed_sql), (
|
|
"Explicit DROP is unnecessary with ON COMMIT DROP"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_uses_set_local_for_ef_search(self, mock_conn):
|
|
"""hnsw.ef_search must be set with SET LOCAL so the change is scoped
|
|
to the transaction. Without SET LOCAL, the setting would leak onto
|
|
the pooled backend and affect subsequent recall queries that land
|
|
on the same backend."""
|
|
emb = [0.1] * 384
|
|
await compute_semantic_links_ann(
|
|
conn=mock_conn,
|
|
bank_id="bank-1",
|
|
unit_ids=["u1"],
|
|
embeddings=[emb],
|
|
fact_types=["world"],
|
|
)
|
|
|
|
executed_sql = [call.args[0] for call in mock_conn.execute.call_args_list]
|
|
ef_statements = [s for s in executed_sql if "hnsw.ef_search" in s]
|
|
assert ef_statements, "ef_search must be tuned down for retain ANN"
|
|
for stmt in ef_statements:
|
|
assert stmt.strip().startswith("SET LOCAL"), (
|
|
f"hnsw.ef_search must use SET LOCAL, got: {stmt}"
|
|
)
|
|
# And there must not be a RESET — SET LOCAL handles it at commit.
|
|
assert not any("RESET hnsw.ef_search" in s for s in executed_sql)
|