Skip to content

Commit

Permalink
Don't use in-memory account info if it would match account info on disk
Browse files Browse the repository at this point in the history
  • Loading branch information
agoncharov-reef committed Mar 25, 2024
1 parent 86d7574 commit 41cb0f7
Show file tree
Hide file tree
Showing 2 changed files with 60 additions and 11 deletions.
32 changes: 21 additions & 11 deletions b2/_internal/console_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -4236,9 +4236,10 @@ def run_command(self, argv):
logger.info('starting command [%s] with arguments: %s', command, argv)

try:
auth_ret = self.authorize_from_env(command_class)
if auth_ret:
return auth_ret
if command_class.REQUIRES_AUTH:
auth_ret = self.authorize_from_env()
if auth_ret:
return auth_ret
return command.run(args)
except MissingAccountData as e:
logger.exception('ConsoleTool missing account data error')
Expand All @@ -4261,16 +4262,25 @@ def run_command(self, argv):

@classmethod
def _initialize_b2_api(cls, args: argparse.Namespace, kwargs: dict) -> B2Api:
if all(get_keyid_and_key_from_env_vars()
) and args.command_class not in (AuthorizeAccount, ClearAccount):
# when user specifies keys via env variables, we switch to in-memory account info
return _get_inmemory_b2api(**kwargs)

return _get_b2api_for_profile(profile=args.profile, **kwargs)
# here we initialize regular b2 api on disk, and check whether it matches
# the keys from env vars; if they indeed match then there's no need to
# initialize in-memory account info cause it's already stored on disk
b2_api = _get_b2api_for_profile(profile=args.profile, **kwargs)

def authorize_from_env(self, command_class):
if not command_class.REQUIRES_AUTH:
return 0
key_id, key = get_keyid_and_key_from_env_vars()
if key_id and key:
realm = os.environ.get(B2_ENVIRONMENT_ENV_VAR) or 'production'
if (
not b2_api.account_info.is_same_key(key_id, realm) and
args.command_class not in (AuthorizeAccount, ClearAccount)
):
# when user specifies keys via env variables, we switch to in-memory account info
return _get_inmemory_b2api(**kwargs)

return b2_api

def authorize_from_env(self) -> int:

key_id, key = get_keyid_and_key_from_env_vars()

Expand Down
39 changes: 39 additions & 0 deletions test/integration/test_b2_command_line.py
Original file line number Diff line number Diff line change
Expand Up @@ -216,6 +216,45 @@ def test_command_with_env_vars_not_saving_credentials(
account_info.get_application_key_id()


@pytest.mark.apiver(from_ver=4)
def test_command_with_env_vars_reusing_existing_account_info(
b2_tool,
application_key,
application_key_id,
account_info_file,
bucket_name,
b2_uri_args,
):
"""
When calling any command with credentials passed via env vars, and the account
info file already contains the same credentials, we want to use filesystem for
storing cache, not the in-memory cache.
"""

assert B2_APPLICATION_KEY_ID_ENV_VAR not in os.environ
assert B2_APPLICATION_KEY_ENV_VAR not in os.environ

assert account_info_file.exists()
account_info = SqliteAccountInfo()
assert account_info.get_application_key() == application_key
assert account_info.get_application_key_id() == application_key_id

account_info.remove_bucket_name(bucket_name)
assert account_info.get_bucket_id_or_none_from_bucket_name(bucket_name) is None

b2_tool.should_succeed(
['ls', '--long', *b2_uri_args(bucket_name)],
additional_env={
B2_APPLICATION_KEY_ID_ENV_VAR: application_key_id,
B2_APPLICATION_KEY_ENV_VAR: application_key,
}
)

assert account_info_file.exists()
account_info = SqliteAccountInfo()
assert account_info.get_bucket_id_or_none_from_bucket_name(bucket_name) is not None


@pytest.fixture
def uploaded_sample_file(b2_tool, bucket_name, sample_filepath):
return b2_tool.should_succeed_json(
Expand Down

0 comments on commit 41cb0f7

Please sign in to comment.