From bc65fab88ae4731fb937d0b0e280348222676077 Mon Sep 17 00:00:00 2001
From: "google-labs-jules[bot]"
 <161369871+google-labs-jules[bot]@users.noreply.github.com>
Date: Sun, 19 Jul 2026 07:44:42 +0000
Subject: [PATCH] perf: optimize table existence checks in OMOP loaders to
 avoid N+1 schema queries

This change refactors `validate_omop_database` and `_upsert_sql_tables` to
query the database schema for table names exactly once for supported
databases (SQLite and DuckDB) using `sqlite_master`. It then uses this
pre-fetched set to quickly check for table existence, bypassing the
slower and repetitive `SELECT 1 FROM table LIMIT 1` exception-driven check.

This significantly improves the performance of validation and upsert
operations by avoiding N+1 schema queries when repeatedly checking for
tables, while still falling back to the original safe check when
necessary (e.g. for non-supported database dialects).

Co-authored-by: zrt219 <199104500+zrt219@users.noreply.github.com>
---
 openmed/openmed/interop/omop/cdm_loader.py | 27 +++++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/openmed/openmed/interop/omop/cdm_loader.py b/openmed/openmed/interop/omop/cdm_loader.py
index 6c1d803..d012e2e 100644
--- a/openmed/openmed/interop/omop/cdm_loader.py
+++ b/openmed/openmed/interop/omop/cdm_loader.py
@@ -846,10 +846,15 @@ def validate_omop_tables(
 def validate_omop_database(con: Any) -> tuple[OmopConstraintViolation, ...]:
     """Validate persisted OMOP tables using the same PHI-free violation shape."""
 
+    existing_tables = _database_tables(con)
     tables = {
         table: tuple(_select_all(con, table))
         for table in _TABLE_ORDER
-        if _table_exists(con, table)
+        if (
+            table in existing_tables
+            if existing_tables is not None
+            else _table_exists(con, table)
+        )
     }
     return validate_omop_tables(
         OmopCdmTables(
@@ -1181,10 +1186,15 @@ def _ordered_row(table: str, row: Mapping[str, Any]) -> dict[str, Any]:
 
 
 def _upsert_sql_tables(con: Any, tables: OmopCdmTables) -> None:
+    existing_tables = _database_tables(con)
     for table in _TABLE_ORDER:
         rows = tables.table(table)
         if not rows:
             continue
+        if existing_tables is not None and table not in existing_tables:
+            continue
+        if existing_tables is None and not _table_exists(con, table):
+            continue
         columns = _SCHEMA_COLUMNS[table]
         column_sql = ", ".join(_quote_identifier(column) for column in columns)
         placeholder_sql = ", ".join("?" for _ in columns)
@@ -1237,6 +1247,21 @@ def _select_all(con: Any, table: str) -> list[dict[str, Any]]:
     return [dict(zip(columns, row)) for row in rows]
 
 
+def _database_tables(con: Any) -> frozenset[str] | None:
+    # Only use fast path for supported in-memory or file-based DBs
+    # where sqlite_master is safe and does not poison transactions.
+    dialect = type(con).__module__.split(".")[0]
+    if dialect in {"sqlite3", "duckdb"}:
+        try:
+            rows = con.execute(
+                "SELECT name FROM sqlite_master WHERE type IN ('table', 'view')"
+            ).fetchall()
+            return frozenset(r[0] for r in rows)
+        except Exception:
+            return None
+    return None
+
+
 def _table_exists(con: Any, table: str) -> bool:
     try:
         con.execute(f"SELECT 1 FROM {table} LIMIT 1")