Merge pull request #314 from reef-technologies/sqlite-account-info-profiles

Allow SqliteAccountInfo to use different profile files

Author: ppolewicz

CHANGELOG.md

@@ -22,6 +22,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Add an option to set a custom file version class to FileVersionFactory
 * Add an option for B2Api to turn off hash checking for downloaded files
 * Add an option for B2Api to set write buffer size for DownloadedFile.save_to method
+* Add support for multiple profile files for SqliteAccountInfo
 
 ### Fixed
 * Fix copying objects larger than 1TB

b2sdk/_v3/__init__.py

@@ -34,7 +34,10 @@
 from b2sdk.account_info.abstract import AbstractAccountInfo
 from b2sdk.account_info.in_memory import InMemoryAccountInfo
 from b2sdk.account_info.sqlite_account_info import SqliteAccountInfo
-from b2sdk.account_info.sqlite_account_info import B2_ACCOUNT_INFO_ENV_VAR, B2_ACCOUNT_INFO_DEFAULT_FILE, XDG_CONFIG_HOME_ENV_VAR
+from b2sdk.account_info.sqlite_account_info import B2_ACCOUNT_INFO_ENV_VAR
+from b2sdk.account_info.sqlite_account_info import B2_ACCOUNT_INFO_DEFAULT_FILE
+from b2sdk.account_info.sqlite_account_info import B2_ACCOUNT_INFO_PROFILE_FILE
+from b2sdk.account_info.sqlite_account_info import XDG_CONFIG_HOME_ENV_VAR
 from b2sdk.account_info.stub import StubAccountInfo
 from b2sdk.account_info.upload_url_pool import UploadUrlPool
 from b2sdk.account_info.upload_url_pool import UrlPoolAccountInfo

b2sdk/account_info/sqlite_account_info.py

@@ -11,19 +11,22 @@
 import json
 import logging
 import os
+import re
+import sqlite3
 import stat
 import threading
-from typing import Optional, List
 
-from .exception import (CorruptAccountInfo, MissingAccountData)
-from .upload_url_pool import UrlPoolAccountInfo
+from typing import List, Optional
 
-import sqlite3
+from .exception import CorruptAccountInfo, MissingAccountData
+from .upload_url_pool import UrlPoolAccountInfo
 
 logger = logging.getLogger(__name__)
 
 B2_ACCOUNT_INFO_ENV_VAR = 'B2_ACCOUNT_INFO'
-B2_ACCOUNT_INFO_DEFAULT_FILE = '~/.b2_account_info'
+B2_ACCOUNT_INFO_DEFAULT_FILE = os.path.join('~', '.b2_account_info')
+B2_ACCOUNT_INFO_PROFILE_FILE = os.path.join('~', '.b2db-{profile}.sqlite')
+B2_ACCOUNT_INFO_PROFILE_NAME_REGEXP = re.compile(r'[a-zA-Z0-9_\-]{1,64}')
 XDG_CONFIG_HOME_ENV_VAR = 'XDG_CONFIG_HOME'
 
 DEFAULT_ABSOLUTE_MINIMUM_PART_SIZE = 5000000  # this value is used ONLY in migrating db, and in v1 wrapper, it is not
@@ -39,7 +42,7 @@ class SqliteAccountInfo(UrlPoolAccountInfo):
     completed.
     """
 
-    def __init__(self, file_name=None, last_upgrade_to_run=None):
+    def __init__(self, file_name=None, last_upgrade_to_run=None, profile: Optional[str] = None):
         """
         Initialize SqliteAccountInfo.
 
@@ -49,6 +52,11 @@ def __init__(self, file_name=None, last_upgrade_to_run=None):
 
         SqliteAccountInfo currently checks locations in the following order:
 
+        If ``profile`` arg is provided:
+        * ``{XDG_CONFIG_HOME_ENV_VAR}/b2/db-<profile>.sqlite``, if ``{XDG_CONFIG_HOME_ENV_VAR}`` env var is set
+        * ``{B2_ACCOUNT_INFO_PROFILE_FILE}``
+
+        Otherwise:
         * ``file_name``, if truthy
         * ``{B2_ACCOUNT_INFO_ENV_VAR}`` env var's value, if set
         * ``{B2_ACCOUNT_INFO_DEFAULT_FILE}``, if it exists
@@ -62,21 +70,7 @@ def __init__(self, file_name=None, last_upgrade_to_run=None):
         """
         self.thread_local = threading.local()
 
-        if file_name:
-            user_account_info_path = file_name
-        elif B2_ACCOUNT_INFO_ENV_VAR in os.environ:
-            user_account_info_path = os.environ[B2_ACCOUNT_INFO_ENV_VAR]
-        elif os.path.exists(os.path.expanduser(B2_ACCOUNT_INFO_DEFAULT_FILE)):
-            user_account_info_path = B2_ACCOUNT_INFO_DEFAULT_FILE
-        elif XDG_CONFIG_HOME_ENV_VAR in os.environ:
-            config_home = os.environ[XDG_CONFIG_HOME_ENV_VAR]
-            user_account_info_path = os.path.join(config_home, 'b2', 'account_info')
-            if not os.path.exists(os.path.join(config_home, 'b2')):
-                os.makedirs(os.path.join(config_home, 'b2'), mode=0o755)
-        else:
-            user_account_info_path = B2_ACCOUNT_INFO_DEFAULT_FILE
-
-        self.filename = os.path.expanduser(user_account_info_path)
+        self.filename = self._get_user_account_info_path(file_name=file_name, profile=profile)
         logger.debug('%s file path to use: %s', self.__class__.__name__, self.filename)
 
         self._validate_database(last_upgrade_to_run)
@@ -90,10 +84,42 @@ def __init__(self, file_name=None, last_upgrade_to_run=None):
             **dict(
                 B2_ACCOUNT_INFO_ENV_VAR=B2_ACCOUNT_INFO_ENV_VAR,
                 B2_ACCOUNT_INFO_DEFAULT_FILE=B2_ACCOUNT_INFO_DEFAULT_FILE,
+                B2_ACCOUNT_INFO_PROFILE_FILE=B2_ACCOUNT_INFO_PROFILE_FILE,
                 XDG_CONFIG_HOME_ENV_VAR=XDG_CONFIG_HOME_ENV_VAR,
             )
         )
 
+    @classmethod
+    def _get_user_account_info_path(cls, file_name=None, profile=None):
+        if profile and not B2_ACCOUNT_INFO_PROFILE_NAME_REGEXP.match(profile):
+            raise ValueError('Invalid profile name: {}'.format(profile))
+
+        if file_name:
+            if profile:
+                raise ValueError('Provide either file_name or profile, not both')
+            user_account_info_path = file_name
+        elif B2_ACCOUNT_INFO_ENV_VAR in os.environ:
+            if profile:
+                raise ValueError(
+                    'Provide either {} env var or profile, not both'.
+                    format(B2_ACCOUNT_INFO_ENV_VAR)
+                )
+            user_account_info_path = os.environ[B2_ACCOUNT_INFO_ENV_VAR]
+        elif os.path.exists(os.path.expanduser(B2_ACCOUNT_INFO_DEFAULT_FILE)) and not profile:
+            user_account_info_path = B2_ACCOUNT_INFO_DEFAULT_FILE
+        elif XDG_CONFIG_HOME_ENV_VAR in os.environ:
+            config_home = os.environ[XDG_CONFIG_HOME_ENV_VAR]
+            file_name = 'db-{}.sqlite'.format(profile) if profile else 'account_info'
+            user_account_info_path = os.path.join(config_home, 'b2', file_name)
+            if not os.path.exists(os.path.join(config_home, 'b2')):
+                os.makedirs(os.path.join(config_home, 'b2'), mode=0o755)
+        elif profile:
+            user_account_info_path = B2_ACCOUNT_INFO_PROFILE_FILE.format(profile=profile)
+        else:
+            user_account_info_path = B2_ACCOUNT_INFO_DEFAULT_FILE
+
+        return os.path.expanduser(user_account_info_path)
+
     def _validate_database(self, last_upgrade_to_run=None):
         """
         Make sure that the database is openable.  Removes the file if it's not.

noxfile.py

@@ -30,6 +30,7 @@
     "pytest-cov==3.0.0",
     "pytest-mock==3.6.1",
     'pytest-lazy-fixture==0.6.3',
+    'pyfakefs==4.5.6',
 ]
 REQUIREMENTS_BUILD = ['setuptools>=20.2']
 
@@ -107,7 +108,7 @@ def unit(session):
     """Run unit tests."""
     install_myself(session)
     session.install(*REQUIREMENTS_TEST)
-    args = ['--doctest-modules']
+    args = ['--doctest-modules', '-p', 'pyfakefs']
     if not SKIP_COVERAGE:
         args += ['--cov=b2sdk', '--cov-branch', '--cov-report=xml']
     # TODO: Use session.parametrize for apiver

test/unit/account_info/test_sqlite_account_info.py

@@ -8,9 +8,18 @@
 #
 ######################################################################
 
+import os
+import unittest.mock as mock
+
 import pytest
 
-from apiver_deps import AbstractAccountInfo
+from apiver_deps import (
+    B2_ACCOUNT_INFO_DEFAULT_FILE,
+    B2_ACCOUNT_INFO_ENV_VAR,
+    XDG_CONFIG_HOME_ENV_VAR,
+    AbstractAccountInfo,
+    SqliteAccountInfo,
+)
 
 from .fixtures import *
 
@@ -55,3 +64,67 @@ def test_migrate_to_4(self):
                 "SELECT recommended_part_size, absolute_minimum_part_size from account"
             ).fetchone()
         assert (100, 5000000) == sizes
+
+
+class TestSqliteAccountProfileFileLocation:
+    @pytest.fixture(autouse=True)
+    def setup(self, monkeypatch, fs):
+        monkeypatch.delenv(B2_ACCOUNT_INFO_ENV_VAR, raising=False)
+        monkeypatch.delenv(XDG_CONFIG_HOME_ENV_VAR, raising=False)
+
+    def test_invalid_profile_name(self):
+        with pytest.raises(ValueError):
+            SqliteAccountInfo._get_user_account_info_path(profile='&@(*$')
+
+    def test_profile_and_file_name_conflict(self):
+        with pytest.raises(ValueError):
+            SqliteAccountInfo._get_user_account_info_path(file_name='foo', profile='bar')
+
+    def test_profile_and_env_var_conflict(self, monkeypatch):
+        monkeypatch.setenv(B2_ACCOUNT_INFO_ENV_VAR, 'foo')
+        with pytest.raises(ValueError):
+            SqliteAccountInfo._get_user_account_info_path(profile='bar')
+
+    def test_profile_and_xdg_config_env_var(self, monkeypatch):
+        monkeypatch.setenv(XDG_CONFIG_HOME_ENV_VAR, os.path.join('~', 'custom'))
+        account_info_path = SqliteAccountInfo._get_user_account_info_path(profile='secondary')
+        assert account_info_path == os.path.expanduser(
+            os.path.join('~', 'custom', 'b2', 'db-secondary.sqlite')
+        )
+
+    def test_profile(self):
+        account_info_path = SqliteAccountInfo._get_user_account_info_path(profile='foo')
+        assert account_info_path == os.path.expanduser(os.path.join('~', '.b2db-foo.sqlite'))
+
+    def test_file_name(self):
+        account_info_path = SqliteAccountInfo._get_user_account_info_path(
+            file_name=os.path.join('~', 'foo')
+        )
+        assert account_info_path == os.path.expanduser(os.path.join('~', 'foo'))
+
+    def test_env_var(self, monkeypatch):
+        monkeypatch.setenv(B2_ACCOUNT_INFO_ENV_VAR, os.path.join('~', 'foo'))
+        account_info_path = SqliteAccountInfo._get_user_account_info_path()
+        assert account_info_path == os.path.expanduser(os.path.join('~', 'foo'))
+
+    def test_default_file_if_exists(self, monkeypatch):
+        # ensure that XDG_CONFIG_HOME_ENV_VAR doesn't matter if default file exists
+        monkeypatch.setenv(XDG_CONFIG_HOME_ENV_VAR, 'some')
+        account_file_path = os.path.expanduser(B2_ACCOUNT_INFO_DEFAULT_FILE)
+        parent_dir = os.path.abspath(os.path.join(account_file_path, os.pardir))
+        os.makedirs(parent_dir, exist_ok=True)
+        with open(account_file_path, 'w') as account_file:
+            account_file.write('')
+        account_info_path = SqliteAccountInfo._get_user_account_info_path()
+        assert account_info_path == os.path.expanduser(B2_ACCOUNT_INFO_DEFAULT_FILE)
+
+    def test_xdg_config_env_var(self, monkeypatch):
+        monkeypatch.setenv(XDG_CONFIG_HOME_ENV_VAR, os.path.join('~', 'custom'))
+        account_info_path = SqliteAccountInfo._get_user_account_info_path()
+        assert account_info_path == os.path.expanduser(
+            os.path.join('~', 'custom', 'b2', 'account_info')
+        )
+
+    def test_default_file(self):
+        account_info_path = SqliteAccountInfo._get_user_account_info_path()
+        assert account_info_path == os.path.expanduser(B2_ACCOUNT_INFO_DEFAULT_FILE)