Address account info bw compatibility issue caused by multi-bucket keys

Author: olzhasar-reef

b2sdk/_internal/account_info/sqlite_account_info.py

@@ -280,7 +280,7 @@ def _create_tables(self, conn, last_upgrade_to_run):
         """
         )
         # By default, we run all the upgrades
-        last_upgrade_to_run = 5 if last_upgrade_to_run is None else last_upgrade_to_run
+        last_upgrade_to_run = 6 if last_upgrade_to_run is None else last_upgrade_to_run
         # Add the 'allowed' column if it hasn't been yet.
         if 1 <= last_upgrade_to_run:
             self._ensure_update(1, ['ALTER TABLE account ADD COLUMN allowed TEXT;'])
@@ -385,13 +385,17 @@ def _create_tables(self, conn, last_upgrade_to_run):
             )
 
         if 5 <= last_upgrade_to_run:
-            self._migrate_allowed_to_multi_bucket()
+            self._migrate_allowed_to_multi_bucket(version=5)
 
-    def _migrate_allowed_to_multi_bucket(self):
+        # There were reported cases of users with db schema upgraded to version 5 and still containing old single-bucket style allowed dict. This was possible because users of sdk apiver <= v2 still received old-styled allowed dicts (from v3 remote API) and there was no proper transformation mechanism when setting the auth data. The bug has now been addressed, but we need an additional migration to address aforementioned cases.
+        if 6 <= last_upgrade_to_run:
+            self._migrate_allowed_to_multi_bucket(version=6)
+
+    def _migrate_allowed_to_multi_bucket(self, *, version: int):
         """
         Migrate existing allowed json dict to a new multi-bucket keys format
         """
-        if self._get_update_count(5) > 0:
+        if self._get_update_count(version) > 0:
             return
 
         try:
@@ -400,11 +404,15 @@ def _migrate_allowed_to_multi_bucket(self):
             allowed_json = None
 
         if allowed_json is None:
-            self._perform_update(5, [])
+            self._perform_update(version, [])
             return
 
         allowed = json.loads(allowed_json)
 
+        if 'buckets' in allowed:
+            self._perform_update(version, [])
+            return
+
         bucket_id = allowed.pop('bucketId')
         bucket_name = allowed.pop('bucketName')
 
@@ -416,7 +424,7 @@ def _migrate_allowed_to_multi_bucket(self):
         allowed_text = json.dumps(allowed)
         stmt = f"UPDATE account SET allowed = ('{allowed_text}');"
 
-        self._perform_update(5, [stmt])
+        self._perform_update(version, [stmt])
 
     def _ensure_update(self, update_number, update_commands: list[str]):
         """
@@ -469,7 +477,6 @@ def _set_auth_data(
         allowed,
         application_key_id,
     ):
-        assert self.allowed_is_valid(allowed)
         with self._get_connection() as conn:
             conn.execute('DELETE FROM account;')
             conn.execute('DELETE FROM bucket;')

b2sdk/v2/account_info.py

@@ -9,8 +9,8 @@
 ######################################################################
 from __future__ import annotations
 
-import json
 
+from copy import copy
 from b2sdk import v3
 from .exception import MissingAccountData
 
@@ -33,49 +33,22 @@ def allowed_is_valid(cls, allowed):
             and ('namePrefix' in allowed)
         )
 
+    @classmethod
+    def _convert_allowed_single_to_multi_bucket(cls, allowed):
+        new = copy(allowed)
 
-class AbstractAccountInfo(_OldAllowedMixin, v3.AbstractAccountInfo):
-    def list_bucket_names_ids(self):
-        return []  # Removed @abstractmethod decorator
-
-
-class UrlPoolAccountInfo(_OldAllowedMixin, v3.UrlPoolAccountInfo):
-    pass
+        bucket_id = new.pop('bucketId')
+        bucket_name = new.pop('bucketName')
 
+        if bucket_id is not None:
+            new['buckets'] = [{'id': bucket_id, 'name': bucket_name}]
+        else:
+            new['buckets'] = None
 
-class InMemoryAccountInfo(_OldAllowedMixin, v3.InMemoryAccountInfo):
-    pass
+        return new
 
-
-class SqliteAccountInfo(_OldAllowedMixin, v3.SqliteAccountInfo):
     def get_allowed(self):
-        """
-        Return 'allowed' dictionary info.
-        Example:
-
-        .. code-block:: python
-
-            {
-                "bucketId": null,
-                "bucketName": null,
-                "capabilities": [
-                    "listKeys",
-                    "writeKeys"
-                ],
-                "namePrefix": null
-            }
-
-        The 'allowed' column was not in the original schema, so it may be NULL.
-
-        :rtype: dict
-        """
-        allowed_json = self._get_account_info_or_raise('allowed')
-        if allowed_json is None:
-            return self.DEFAULT_ALLOWED
-
-        allowed = json.loads(allowed_json)
-
-        # convert a multi-bucket key to a single bucket
+        allowed = super().get_allowed()
 
         if 'buckets' in allowed:
             buckets = allowed.pop('buckets')
@@ -89,6 +62,53 @@ def get_allowed(self):
 
         return allowed
 
+    def _set_auth_data(
+        self,
+        account_id,
+        auth_token,
+        api_url,
+        download_url,
+        recommended_part_size,
+        absolute_minimum_part_size,
+        application_key,
+        realm,
+        s3_api_url,
+        allowed,
+        application_key_id,
+    ):
+        new_allowed = self._convert_allowed_single_to_multi_bucket(allowed)
+
+        super()._set_auth_data(
+            account_id,
+            auth_token,
+            api_url,
+            download_url,
+            recommended_part_size,
+            absolute_minimum_part_size,
+            application_key,
+            realm,
+            s3_api_url,
+            new_allowed,
+            application_key_id,
+        )
+
+
+class AbstractAccountInfo(_OldAllowedMixin, v3.AbstractAccountInfo):
+    def list_bucket_names_ids(self):
+        return []  # Removed @abstractmethod decorator
+
+
+class UrlPoolAccountInfo(_OldAllowedMixin, v3.UrlPoolAccountInfo):
+    pass
+
+
+class InMemoryAccountInfo(_OldAllowedMixin, v3.InMemoryAccountInfo):
+    pass
+
+
+class SqliteAccountInfo(_OldAllowedMixin, v3.SqliteAccountInfo):
+    pass
+
 
 class StubAccountInfo(_OldAllowedMixin, v3.StubAccountInfo):
     pass

changelog.d/+multi-bucket-keys-account-info-bw-compat.fixed.md

@@ -0,0 +1 @@
+Address backwards compatibility issue for sqlite account info caused by the migration of schema to a new multi-bucket format.

test/unit/account_info/fixtures.py

@@ -69,13 +69,9 @@ 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, schema_0=False):
+    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'))
-        if schema_0:
-            last_upgrade_to_run = 0
-        else:
-            last_upgrade_to_run = None
         return SqliteAccountInfo(file_name, last_upgrade_to_run)
 
     return get_account_info

test/unit/account_info/test_sqlite_account_info.py

@@ -9,6 +9,7 @@
 ######################################################################
 from __future__ import annotations
 
+import json
 import os
 
 import pytest
@@ -32,30 +33,30 @@ def setup(self, sqlite_account_info_factory, account_info_default_data_schema_0)
 
     def test_upgrade_1_default_allowed(self):
         """The 'allowed' field should be the default for upgraded databases."""
-        old_account_info = self.sqlite_account_info_factory(schema_0=True)
+        old_account_info = self.sqlite_account_info_factory(last_upgrade_to_run=0)
         old_account_info.set_auth_data_with_schema_0_for_test(**self.account_info_default_data)
         new_account_info = self.sqlite_account_info_factory(file_name=old_account_info.filename)
 
         assert AbstractAccountInfo.DEFAULT_ALLOWED == new_account_info.get_allowed()
 
     def test_upgrade_2_default_app_key(self):
         """The 'application_key_id' field should default to the account ID."""
-        old_account_info = self.sqlite_account_info_factory(schema_0=True)
+        old_account_info = self.sqlite_account_info_factory(last_upgrade_to_run=0)
         old_account_info.set_auth_data_with_schema_0_for_test(**self.account_info_default_data)
         new_account_info = self.sqlite_account_info_factory(file_name=old_account_info.filename)
 
         assert 'account_id' == new_account_info.get_application_key_id()
 
     def test_upgrade_3_default_s3_api_url(self):
         """The 's3_api_url' field should be set."""
-        old_account_info = self.sqlite_account_info_factory(schema_0=True)
+        old_account_info = self.sqlite_account_info_factory(last_upgrade_to_run=0)
         old_account_info.set_auth_data_with_schema_0_for_test(**self.account_info_default_data)
         new_account_info = self.sqlite_account_info_factory(file_name=old_account_info.filename)
 
         assert '' == new_account_info.get_s3_api_url()
 
     def test_migrate_to_4(self):
-        old_account_info = self.sqlite_account_info_factory(schema_0=True)
+        old_account_info = self.sqlite_account_info_factory(last_upgrade_to_run=0)
         old_account_info.set_auth_data_with_schema_0_for_test(**self.account_info_default_data)
         new_account_info = self.sqlite_account_info_factory(file_name=old_account_info.filename)
 
@@ -82,19 +83,48 @@ def test_migrate_to_4(self):
             ),
         ],
     )
-    @pytest.mark.apiver(2)
-    def test_migrate_to_5_v2(
-        self, v2_sqlite_account_info_factory, account_info_default_data, allowed
+    @pytest.mark.apiver(to_ver=2)
+    @pytest.mark.parametrize(
+        ('start_version', 'end_version'),
+        [
+            (4, 5),
+            (5, 6),
+        ],
+    )
+    def test_migrate_to_5_prior_to_v3(
+        self,
+        sqlite_account_info_factory,
+        account_info_default_data: dict,
+        allowed: dict,
+        start_version: int,
+        end_version: int,
     ):
-        old_account_info = v2_sqlite_account_info_factory()
-        account_info_default_data.update({'allowed': allowed})
+        old_account_info = sqlite_account_info_factory(last_upgrade_to_run=start_version)
         old_account_info.set_auth_data(**account_info_default_data)
+
+        allowed_text = json.dumps(allowed)
+        with old_account_info._get_connection() as conn:
+            conn.execute(f"UPDATE account SET allowed = ('{allowed_text}');")
+
+        assert old_account_info._get_update_count(start_version) == 1
+        assert old_account_info._get_update_count(end_version) == 0
+
         assert old_account_info.get_allowed() == allowed
 
-        new_account_info = self.sqlite_account_info_factory(file_name=old_account_info.filename)
+        new_account_info = self.sqlite_account_info_factory(
+            file_name=old_account_info.filename, last_upgrade_to_run=end_version
+        )
 
+        assert new_account_info._get_update_count(end_version) == 1
         assert new_account_info.get_allowed() == allowed
 
+        with new_account_info._get_connection() as conn:
+            new_allowed = json.loads(conn.execute('SELECT allowed FROM account;').fetchone()[0])
+
+        assert 'buckets' in new_allowed
+        assert 'bucketId' not in new_allowed
+        assert 'bucketName' not in new_allowed
+
     @pytest.mark.parametrize(
         ('old_allowed', 'exp_allowed'),
         [
@@ -126,22 +156,40 @@ def test_migrate_to_5_v2(
             ),
         ],
     )
+    @pytest.mark.parametrize(
+        ('start_version', 'end_version'),
+        [
+            (4, 5),
+            (5, 6),
+        ],
+    )
     @pytest.mark.apiver(from_ver=3)
     def test_migrate_to_5_v3(
         self,
-        v2_sqlite_account_info_factory,
-        account_info_default_data,
-        old_allowed,
-        exp_allowed,
+        sqlite_account_info_factory,
+        account_info_default_data: dict,
+        old_allowed: dict,
+        exp_allowed: dict,
+        start_version: int,
+        end_version: int,
     ):
-        old_account_info = v2_sqlite_account_info_factory()
-        account_info_default_data.update({'allowed': old_allowed})
+        old_account_info = sqlite_account_info_factory(last_upgrade_to_run=start_version)
         old_account_info.set_auth_data(**account_info_default_data)
+
+        allowed_text = json.dumps(old_allowed)
+        with old_account_info._get_connection() as conn:
+            conn.execute(f"UPDATE account SET allowed = ('{allowed_text}');")
+
         assert old_account_info.get_allowed() == old_allowed
+        assert old_account_info._get_update_count(start_version) == 1
+        assert old_account_info._get_update_count(end_version) == 0
 
-        new_account_info = self.sqlite_account_info_factory(file_name=old_account_info.filename)
+        new_account_info = self.sqlite_account_info_factory(
+            file_name=old_account_info.filename, last_upgrade_to_run=end_version
+        )
 
         assert new_account_info.get_allowed() == exp_allowed
+        assert new_account_info._get_update_count(end_version) == 1
 
     @pytest.mark.apiver(2)
     def test_multi_bucket_key_error_apiver_v2(