Merge pull request #566 from Backblaze/close-sqlite-connections

Close sqlite connections in SqliteAccountInfo

Author: olzhasar-reef

b2sdk/_internal/account_info/sqlite_account_info.py

@@ -17,6 +17,7 @@
 import stat
 import sys
 import threading
+from contextlib import contextmanager
 
 from .exception import CorruptAccountInfo, MissingAccountData
 from .upload_url_pool import UrlPoolAccountInfo
@@ -152,7 +153,7 @@ def _validate_database(self, last_upgrade_to_run=None):
 
         # If we can connect to the database, and do anything, then all is good.
         try:
-            with self._connect() as conn:
+            with self._temp_connection() as conn:
                 self._create_tables(conn, last_upgrade_to_run)
                 return
         except sqlite3.DatabaseError:
@@ -178,7 +179,7 @@ def _validate_database(self, last_upgrade_to_run=None):
                 # create a database
                 self._create_database(last_upgrade_to_run)
                 # add the data from the JSON file
-                with self._connect() as conn:
+                with self._temp_connection() as conn:
                     self._create_tables(conn, last_upgrade_to_run)
                     insert_statement = """
                         INSERT INTO account
@@ -214,6 +215,24 @@ def _get_connection(self):
     def _connect(self):
         return sqlite3.connect(self.filename, isolation_level='EXCLUSIVE')
 
+    @contextmanager
+    def _temp_connection(self):
+        """A one-off sqlite connection for setup purposes, that is not cached in thread_local."""
+        conn = self._connect()
+        try:
+            with conn:
+                yield conn
+        finally:
+            conn.close()
+
+    def close(self):
+        connection = getattr(self.thread_local, 'connection', None)
+        if connection is None:
+            return
+
+        connection.close()
+        del self.thread_local.connection
+
     def _create_database(self, last_upgrade_to_run):
         """
         Make sure that the database is created and has appropriate file permissions.

changelog.d/+sqlite-account-info-connections.fixed.md

@@ -0,0 +1 @@
+Fix `SqliteAccountInfo` to explicitly close temporary sqlite connections used during setup, add a `close()` method for cached connection.

test/integration/test_raw_api.py

@@ -15,7 +15,6 @@
 import re
 import sys
 import time
-import traceback
 from typing import List
 
 import pytest
@@ -70,12 +69,8 @@ def test_raw_api(dont_cleanup_old_buckets):
 
     print()
 
-    try:
-        raw_api = B2RawHTTPApi(B2Http())
-        raw_api_test_helper(raw_api, not dont_cleanup_old_buckets)
-    except Exception:
-        traceback.print_exc(file=sys.stdout)
-        pytest.fail('test_raw_api failed')
+    raw_api = B2RawHTTPApi(B2Http())
+    raw_api_test_helper(raw_api, not dont_cleanup_old_buckets)
 
 
 def authorize_raw_api(raw_api):

test/unit/account_info/fixtures.py

@@ -13,8 +13,7 @@
 from apiver_deps import InMemoryAccountInfo, SqliteAccountInfo
 from pytest_lazy_fixtures import lf
 
-from b2sdk.v2 import SqliteAccountInfo as V2SqliteAccountInfo
-from b2sdk.v3 import SqliteAccountInfo as V3SqliteAccountInfo
+_USE_TMPDIR = object()
 
 
 @pytest.fixture
@@ -69,32 +68,22 @@ def in_memory_account_info(in_memory_account_info_factory):
 
 @pytest.fixture
 def sqlite_account_info_factory(tmpdir):
-    def get_account_info(file_name=None, last_upgrade_to_run: int | None = None):
-        if file_name is None:
-            file_name = str(tmpdir.join('b2_account_info'))
-        return SqliteAccountInfo(file_name, last_upgrade_to_run)
-
-    return get_account_info
+    created_account_infos = []
 
-
-@pytest.fixture
-def v2_sqlite_account_info_factory(tmpdir):
-    def get_account_info(file_name=None):
-        if file_name is None:
+    def get_account_info(
+        file_name=_USE_TMPDIR, *, account_info_class=SqliteAccountInfo, **account_info_kwargs
+    ):
+        if file_name is _USE_TMPDIR:
             file_name = str(tmpdir.join('b2_account_info'))
-        return V2SqliteAccountInfo(file_name, last_upgrade_to_run=4)
 
-    return get_account_info
+        account_info = account_info_class(file_name=file_name, **account_info_kwargs)
+        created_account_infos.append(account_info)
+        return account_info
 
+    yield get_account_info
 
-@pytest.fixture
-def v3_sqlite_account_info_factory(tmpdir):
-    def get_account_info(file_name=None):
-        if file_name is None:
-            file_name = str(tmpdir.join('b2_account_info'))
-        return V3SqliteAccountInfo(file_name)
-
-    return get_account_info
+    for account_info in created_account_infos:
+        account_info.close()
 
 
 @pytest.fixture

test/unit/account_info/test_account_info.py

@@ -349,7 +349,8 @@ class TestSqliteAccountInfo(AccountInfoBase):
     PERSISTENCE = True
 
     @pytest.fixture(autouse=True)
-    def setUp(self, request):
+    def setUp(self, request, sqlite_account_info_factory):
+        self.sqlite_account_info_factory = sqlite_account_info_factory
         self.db_path = tempfile.NamedTemporaryFile(
             prefix=f'tmp_b2_tests_{request.node.name}__', delete=True
         ).name
@@ -378,9 +379,7 @@ def test_permissions(self):
         """
         Test that a new database won't be readable by just any user
         """
-        SqliteAccountInfo(
-            file_name=self.db_path,
-        )
+        self.sqlite_account_info_factory(file_name=self.db_path)
         mode = os.stat(self.db_path).st_mode
         assert stat.filemode(mode) == '-rw-------'
 
@@ -428,7 +427,7 @@ def _make_sqlite_account_info(self, env=None, last_upgrade_to_run=None):
         """
         # Override HOME to ensure hermetic tests
         with mock.patch('os.environ', env or {'HOME': self.test_home}):
-            return SqliteAccountInfo(
+            return self.sqlite_account_info_factory(
                 file_name=self.db_path if not env else None,
                 last_upgrade_to_run=last_upgrade_to_run,
             )

test/unit/account_info/test_sqlite_account_info.py

@@ -22,6 +22,9 @@
 )
 from apiver_deps_exception import MissingAccountData
 
+from b2sdk.v2 import SqliteAccountInfo as V2SqliteAccountInfo
+from b2sdk.v3 import SqliteAccountInfo as V3SqliteAccountInfo
+
 from .fixtures import *
 
 
@@ -194,8 +197,7 @@ def test_migrate_to_5_v3(
     @pytest.mark.apiver(2)
     def test_multi_bucket_key_error_apiver_v2(
         self,
-        v2_sqlite_account_info_factory,
-        v3_sqlite_account_info_factory,
+        sqlite_account_info_factory,
         account_info_default_data,
     ):
         allowed = dict(
@@ -204,13 +206,17 @@ def test_multi_bucket_key_error_apiver_v2(
             namePrefix=None,
         )
 
-        v3_account_info = v3_sqlite_account_info_factory()
+        v3_account_info = sqlite_account_info_factory(account_info_class=V3SqliteAccountInfo)
         account_info_default_data.update({'allowed': allowed})
         v3_account_info.set_auth_data(**account_info_default_data)
 
         assert v3_account_info.get_allowed() == allowed
 
-        v2_account_info = v2_sqlite_account_info_factory(file_name=v3_account_info.filename)
+        v2_account_info = sqlite_account_info_factory(
+            file_name=v3_account_info.filename,
+            account_info_class=V2SqliteAccountInfo,
+            last_upgrade_to_run=4,
+        )
         with pytest.raises(MissingAccountData):
             v2_account_info.get_allowed()