Fix pg_trgm unavailability causing startup crash and silent retain failures (#626) (#649)

On managed PostgreSQL services (e.g. Azure Flexible Server), the pg_trgm
extension may not be available, causing two failures:

1. Migration c1a2b3d4e5f6 crashes on CREATE EXTENSION
2. Even if migration is bypassed, the default 'trigram' entity lookup
   strategy uses the % operator which requires pg_trgm, causing retain
   background tasks to fail silently

Changes:
- Migration now gracefully skips pg_trgm and index creation if the
  extension cannot be loaded
- EntityResolver auto-detects pg_trgm availability on first use and
  falls back to 'full' lookup strategy with a warning log

Co-authored-by: coder999999999 <coder999999999@users.noreply.github.com>
This commit is contained in:
Coder 2026-03-23 04:18:59 -04:00 committed by GitHub
parent 2eb1019da9
commit 365fa3ce50
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 34 additions and 2 deletions

View file

@ -11,6 +11,7 @@ block; see migrations.py for how this is handled safely.
from collections.abc import Sequence
import sqlalchemy as sa
from alembic import context, op
revision: str = "c1a2b3d4e5f6"
@ -25,9 +26,21 @@ def _get_schema_prefix() -> str:
def upgrade() -> None:
# pg_trgm ships with every standard PostgreSQL installation as a contrib module.
# pg_trgm ships with most PostgreSQL installations as a contrib module.
# It enables fast similarity lookups via GIN indexes, used for entity name matching.
op.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm")
# On managed services (e.g. Azure Flexible Server), the extension may not be
# available or may require manual enablement. We gracefully skip the index
# creation if the extension cannot be loaded — the entity resolver will
# auto-detect and fall back to the "full" lookup strategy at runtime. See #626.
conn = op.get_bind()
try:
conn.execute(sa.text("CREATE EXTENSION IF NOT EXISTS pg_trgm"))
except Exception:
# Extension not available (managed Postgres, insufficient privileges, etc.)
# Roll back the failed statement and skip index creation.
conn.execute(sa.text("ROLLBACK"))
conn.execute(sa.text("BEGIN"))
return
schema = _get_schema_prefix()
# GIN index on canonical_name enables sub-millisecond trigram similarity queries

View file

@ -75,6 +75,7 @@ class EntityResolver:
"""
self.pool = pool
self.entity_lookup = entity_lookup
self._pg_trgm_checked = False
# Keyed by asyncio task id so concurrent retain batches never mix their
# pending updates. flush_pending_stats() pops only the calling task's items.
self._pending_stats: dict[int, list[_EntityStat]] = {}
@ -202,6 +203,24 @@ class EntityResolver:
taxonomy_lookup: set[str] | None = None,
) -> list[str]:
if self.entity_lookup == "trigram":
# Auto-detect pg_trgm availability on first call and fall back to
# "full" strategy if the extension is not installed. See #626.
if not self._pg_trgm_checked:
self._pg_trgm_checked = True
has_trgm = await conn.fetchval(
"SELECT EXISTS(SELECT 1 FROM pg_extension WHERE extname = 'pg_trgm')"
)
if not has_trgm:
logger.warning(
"pg_trgm extension is not available — falling back to 'full' "
"entity lookup strategy. Install pg_trgm for faster entity "
"resolution on large banks. See: "
"https://github.com/vectorize-io/hindsight/issues/626"
)
self.entity_lookup = "full"
return await self._resolve_entities_batch_full(
conn, bank_id, entities_data, unit_event_date
)
return await self._resolve_entities_batch_trigram(conn, bank_id, entities_data, unit_event_date)
return await self._resolve_entities_batch_full(conn, bank_id, entities_data, unit_event_date)