Fix possible race condition

Author: Maxteabag

sqlit/ui/mixins/autocomplete.py

@@ -21,6 +21,12 @@ class AutocompleteMixin:
     _schema_cache: dict[str, Any] = {}
     _table_metadata: dict[str, tuple[str, str, str | None]] = {}
 
+    def _run_db_call(self: AppProtocol, fn: Any, *args: Any, **kwargs: Any) -> Any:
+        session = getattr(self, "_session", None)
+        if session is not None:
+            return session.executor.submit(fn, *args, **kwargs).result()
+        return fn(*args, **kwargs)
+
     def _get_word_before_cursor(self, text: str, cursor_pos: int) -> tuple[str, str]:
         """Get the current word being typed and the context keyword before it."""
         if cursor_pos <= 0 or cursor_pos > len(text):
@@ -103,7 +109,9 @@ def work() -> None:
                 column_names = []
             else:
                 try:
-                    columns = adapter.get_columns(connection, actual_table_name, database, schema_name)
+                    columns = self._run_db_call(
+                        adapter.get_columns, connection, actual_table_name, database, schema_name
+                    )
                     column_names = [c.name for c in columns]
                 except Exception:
                     column_names = []
@@ -362,6 +370,12 @@ async def _load_schema_cache_async(self: AppProtocol) -> None:
         }
         table_metadata: dict[str, tuple[str, str, str | None]] = {}
 
+        async def run_db_call(fn: Any, *args: Any, **kwargs: Any) -> Any:
+            session = getattr(self, "_session", None)
+            if session is not None:
+                return await session.executor.run_async(fn, *args, **kwargs)
+            return await asyncio.to_thread(fn, *args, **kwargs)
+
         try:
             # Get database list in thread
             databases: list[str | None]
@@ -370,7 +384,7 @@ async def _load_schema_cache_async(self: AppProtocol) -> None:
                 if db and db.lower() not in ("", "master"):
                     databases = [db]
                 else:
-                    all_dbs = await asyncio.to_thread(adapter.get_databases, connection)
+                    all_dbs = await run_db_call(adapter.get_databases, connection)
                     system_dbs = {"master", "tempdb", "model", "msdb"}
                     databases = [d for d in all_dbs if d.lower() not in system_dbs]
             else:
@@ -379,7 +393,7 @@ async def _load_schema_cache_async(self: AppProtocol) -> None:
             for database in databases:
                 try:
                     # Get tables in thread (NO columns - lazy loaded)
-                    tables = await asyncio.to_thread(adapter.get_tables, connection, database)
+                    tables = await run_db_call(adapter.get_tables, connection, database)
                     for schema_name, table_name in tables:
                         display_name = adapter.format_table_name(schema_name, table_name)
                         schema_cache["tables"].append(display_name)
@@ -392,7 +406,7 @@ async def _load_schema_cache_async(self: AppProtocol) -> None:
                             table_metadata[full_name.lower()] = (schema_name, table_name, database)
 
                     # Get views in thread (NO columns - lazy loaded)
-                    views = await asyncio.to_thread(adapter.get_views, connection, database)
+                    views = await run_db_call(adapter.get_views, connection, database)
                     for schema_name, view_name in views:
                         display_name = adapter.format_table_name(schema_name, view_name)
                         schema_cache["views"].append(display_name)
@@ -405,7 +419,7 @@ async def _load_schema_cache_async(self: AppProtocol) -> None:
                             table_metadata[full_name.lower()] = (schema_name, view_name, database)
 
                     if adapter.supports_stored_procedures:
-                        procedures = await asyncio.to_thread(adapter.get_procedures, connection, database)
+                        procedures = await run_db_call(adapter.get_procedures, connection, database)
                         schema_cache["procedures"].extend(procedures)
 
                 except Exception:

sqlit/ui/mixins/tree.py

@@ -32,6 +32,12 @@
 class TreeMixin:
     """Mixin providing tree/explorer functionality."""
 
+    def _run_db_call(self: AppProtocol, fn: Any, *args: Any, **kwargs: Any) -> Any:
+        session = getattr(self, "_session", None)
+        if session is not None:
+            return session.executor.submit(fn, *args, **kwargs).result()
+        return fn(*args, **kwargs)
+
     def _db_type_badge(self, db_type: str) -> str:
         """Get short badge for database type."""
         return get_badge_label(db_type)
@@ -166,7 +172,7 @@ def get_conn_label(config: Any, connected: Any = False) -> str:
                     dbs_node = active_node.add("Databases")
                     dbs_node.data = FolderNode(folder_type="databases")
 
-                    databases = adapter.get_databases(self.current_connection)
+                    databases = self._run_db_call(adapter.get_databases, self.current_connection)
                     for db_name in databases:
                         db_node = dbs_node.add(escape_markup(db_name))
                         db_node.data = DatabaseNode(name=db_name)
@@ -331,7 +337,7 @@ def work() -> None:
                 else:
                     adapter = self._session.adapter
                     conn = self._session.connection
-                    columns = adapter.get_columns(conn, obj_name, db_name, schema_name)
+                    columns = self._run_db_call(adapter.get_columns, conn, obj_name, db_name, schema_name)
 
                 # Update UI from worker thread
                 self.call_from_thread(self._on_columns_loaded, node, db_name, schema_name, obj_name, columns)
@@ -377,27 +383,39 @@ def work() -> None:
                     conn = self._session.connection
 
                     if folder_type == "tables":
-                        items = [("table", s, t) for s, t in adapter.get_tables(conn, db_name)]
+                        items = [("table", s, t) for s, t in self._run_db_call(adapter.get_tables, conn, db_name)]
                     elif folder_type == "views":
-                        items = [("view", s, v) for s, v in adapter.get_views(conn, db_name)]
+                        items = [("view", s, v) for s, v in self._run_db_call(adapter.get_views, conn, db_name)]
                     elif folder_type == "indexes":
                         if adapter.supports_indexes:
-                            items = [("index", i.name, i.table_name) for i in adapter.get_indexes(conn, db_name)]
+                            items = [
+                                ("index", i.name, i.table_name)
+                                for i in self._run_db_call(adapter.get_indexes, conn, db_name)
+                            ]
                         else:
                             items = []
                     elif folder_type == "triggers":
                         if adapter.supports_triggers:
-                            items = [("trigger", t.name, t.table_name) for t in adapter.get_triggers(conn, db_name)]
+                            items = [
+                                ("trigger", t.name, t.table_name)
+                                for t in self._run_db_call(adapter.get_triggers, conn, db_name)
+                            ]
                         else:
                             items = []
                     elif folder_type == "sequences":
                         if adapter.supports_sequences:
-                            items = [("sequence", s.name, "") for s in adapter.get_sequences(conn, db_name)]
+                            items = [
+                                ("sequence", s.name, "")
+                                for s in self._run_db_call(adapter.get_sequences, conn, db_name)
+                            ]
                         else:
                             items = []
                     elif folder_type == "procedures":
                         if adapter.supports_stored_procedures:
-                            items = [("procedure", "", p) for p in adapter.get_procedures(conn, db_name)]
+                            items = [
+                                ("procedure", "", p)
+                                for p in self._run_db_call(adapter.get_procedures, conn, db_name)
+                            ]
                         else:
                             items = []
                     else: