autoimport: track and close SQLite connections across worker threads

CHANGELOG.md

@@ -1,5 +1,6 @@
 # **Upcoming release**
 
+- #886 Track and close SQLite connections created in worker threads in AutoImport (@mcepl)
 - #895 patchedast Starred and keyword now consumes their syntactically expected and ** (@lieryan)
 - #896 patchedast cleanup and refactoring (@lieryan)
 - #905 Ignore `node_modules` directories by default when scanning project files and source folders. (@marcoeangeli)

docs/library.rst

@@ -853,7 +853,12 @@ by passing memory=False to AutoImport constructor. This default will change in
 the future, if you want to always store the autoimport cache in-memory, then
 you need to explicitly pass memory=True.
 
-It must be closed when done with the ``AutoImport.close()`` method.
+It must be closed when done with the ``AutoImport.close()`` method. When
+sharing an instance across threads, wait for worker threads to finish before
+calling ``close()``. It raises ``RopeError`` rather than closing a connection
+owned by a live worker thread. A worker can close its own connection sooner
+with ``close_thread_connection()``. Connections left by finished workers are
+closed when another thread creates a connection or when ``close()`` runs.
 
 AutoImport can search for a name from both modules and statements you can import from them.
 

rope/contrib/autoimport/sqlite.py

@@ -15,9 +15,10 @@
 from hashlib import sha256
 from itertools import chain
 from pathlib import Path
-from threading import local
+from threading import Lock, Thread, current_thread, local
 from typing import (
     TYPE_CHECKING,
+    Dict,
     Generator,
     Iterable,
     Iterator,
@@ -142,17 +143,21 @@ def __init__(
                 "`AutoImport(memory=True)` explicitly.",
                 DeprecationWarning,
             )
+        self._closed = False
+        self._connections: Dict[Thread, sqlite3.Connection] = {}
+        self._connections_lock = Lock()
+        self._observer: Optional[resourceobserver.ResourceObserver] = None
         self.thread_local = local()
         self.connection = self.create_database_connection(
             project=project,
             memory=memory,
         )
         self._setup_db()
         if observe:
-            observer = resourceobserver.ResourceObserver(
+            self._observer = resourceobserver.ResourceObserver(
                 changed=self._changed, moved=self._moved, removed=self._removed
             )
-            project.add_observer(observer)
+            project.add_observer(self._observer)
 
     @classmethod
     def create_database_connection(
@@ -186,10 +191,41 @@ def calculate_project_hash(data: str) -> str:
             else:
                 project_hash = calculate_project_hash(project.ropefolder.real_path)
             return sqlite3.connect(
-                f"file:rope-{project_hash}:?mode=memory&cache=shared", uri=True
+                f"file:rope-{project_hash}:?mode=memory&cache=shared",
+                uri=True,
+                check_same_thread=False,
             )
         else:
-            return sqlite3.connect(project.ropefolder.pathlib / "autoimport.db")
+            return sqlite3.connect(
+                project.ropefolder.pathlib / "autoimport.db",
+                check_same_thread=False,
+            )
+
+    @staticmethod
+    def _commit_and_close(conn: sqlite3.Connection) -> None:
+        with contextlib.suppress(sqlite3.ProgrammingError, sqlite3.OperationalError):
+            conn.commit()
+        with contextlib.suppress(sqlite3.ProgrammingError, sqlite3.OperationalError):
+            conn.close()
+
+    def _close_finished_thread_connections(self) -> None:
+        """Close connections whose owning threads have exited.
+
+        The caller must hold ``_connections_lock``.
+        """
+        for thread, conn in list(self._connections.items()):
+            if not thread.is_alive():
+                del self._connections[thread]
+                self._commit_and_close(conn)
+
+    def _register_connection(self, conn: sqlite3.Connection) -> None:
+        thread = current_thread()
+        with self._connections_lock:
+            if self._closed:
+                self._commit_and_close(conn)
+                raise exceptions.RopeError("AutoImport instance has been closed")
+            self._close_finished_thread_connections()
+            self._connections[thread] = conn
 
     @property
     def connection(self) -> sqlite3.Connection:
@@ -198,16 +234,31 @@ def connection(self) -> sqlite3.Connection:
 
         This makes sure AutoImport can be shared across threads.
         """
+        if self._closed:
+            raise exceptions.RopeError("AutoImport instance has been closed")
         if not hasattr(self.thread_local, "connection"):
-            self.thread_local.connection = self.create_database_connection(
+            conn = self.create_database_connection(
                 project=self.project,
                 memory=self.memory,
             )
+            self._register_connection(conn)
+            self.thread_local.connection = conn
         return self.thread_local.connection
 
     @connection.setter
     def connection(self, value: sqlite3.Connection):
-        self.thread_local.connection = value
+        thread = current_thread()
+        with self._connections_lock:
+            if self._closed:
+                self._commit_and_close(value)
+                raise exceptions.RopeError("AutoImport instance has been closed")
+            self._close_finished_thread_connections()
+            old_conn = getattr(self.thread_local, "connection", None)
+            if old_conn is not None and old_conn is not value:
+                if self._connections.get(thread) is old_conn:
+                    del self._connections[thread]
+            self._connections[thread] = value
+            self.thread_local.connection = value
 
     def _setup_db(self):
         models.Metadata.create_table(self.connection)
@@ -457,10 +508,55 @@ def update_module(self, module: str):
         self._del_if_exist(module)
         self.generate_modules_cache([module])
 
+    def close_thread_connection(self):
+        """Close the SQLite connection owned by the current thread."""
+        conn = getattr(self.thread_local, "connection", None)
+        if conn is not None:
+            thread = current_thread()
+            with self._connections_lock:
+                if self._connections.get(thread) is conn:
+                    del self._connections[thread]
+            with contextlib.suppress(AttributeError):
+                del self.thread_local.connection
+            self._commit_and_close(conn)
+
     def close(self):
-        """Close the autoimport database."""
-        self.connection.commit()
-        self.connection.close()
+        """Close the database after all worker threads have stopped.
+
+        Raises ``RopeError`` if another live thread still owns a connection.
+        """
+        thread = current_thread()
+        with self._connections_lock:
+            if self._closed:
+                return
+            self._close_finished_thread_connections()
+            active_threads = [
+                owner
+                for owner in self._connections
+                if owner is not thread and owner.is_alive()
+            ]
+            if active_threads:
+                raise exceptions.RopeError(
+                    "Cannot close AutoImport while other threads own connections; "
+                    "wait for them to finish first"
+                )
+            self._closed = True
+            connections = list(self._connections.values())
+            self._connections.clear()
+
+        if self._observer is not None:
+            with contextlib.suppress(Exception):
+                self.project.remove_observer(self._observer)
+            self._observer = None
+
+        for conn in connections:
+            self._commit_and_close(conn)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
 
     def get_name_locations(self, name):
         """Return a list of ``(resource, lineno)`` tuples."""