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:
parent
2eb1019da9
commit
365fa3ce50
2 changed files with 34 additions and 2 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue