Skip to content

Commit

Permalink
Merge pull request #1060 from sirosen/close-sqlite-connections-on-exit
Browse files Browse the repository at this point in the history
Refactor to scope the lifecycle of SQLite storage
  • Loading branch information
sirosen authored Jan 6, 2025
2 parents e7ef1ff + a57d8cf commit d2a8289
Show file tree
Hide file tree
Showing 20 changed files with 497 additions and 424 deletions.
11 changes: 7 additions & 4 deletions src/globus_cli/commands/cli_profile_list.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@

import click

from globus_cli.login_manager import is_client_login, token_storage_adapter
from globus_cli.login_manager import LoginManager, is_client_login
from globus_cli.parsing import command
from globus_cli.termio import Field, display, formatters

Expand All @@ -29,13 +29,16 @@ def _profilestr_to_datadict(s: str) -> dict[str, t.Any] | None:


def _parse_and_filter_profiles(
login_manager: LoginManager,
all: bool,
) -> tuple[list[dict[str, t.Any]], list[dict[str, t.Any]]]:
globus_env = os.getenv("GLOBUS_SDK_ENVIRONMENT", "production")

client_profiles = []
user_profiles = []
for n in token_storage_adapter().iter_namespaces(include_config_namespaces=True):
for n in login_manager.storage.adapter.iter_namespaces(
include_config_namespaces=True
):
data = _profilestr_to_datadict(n)
if not data: # skip any parse failures
continue
Expand Down Expand Up @@ -86,8 +89,8 @@ def cli_profile_list(*, all: bool) -> None:
These are the values for GLOBUS_PROFILE which have been recorded, as well as
GLOBUS_CLI_CLIENT_ID values which have been used.
"""

client_profiles, user_profiles = _parse_and_filter_profiles(all)
login_manager = LoginManager()
client_profiles, user_profiles = _parse_and_filter_profiles(login_manager, all)

if user_profiles:
fields = [
Expand Down
7 changes: 3 additions & 4 deletions src/globus_cli/commands/collection/_common.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import click
import globus_sdk

from globus_cli.login_manager.utils import get_current_identity_id
from globus_cli.login_manager import LoginManager
from globus_cli.termio import Field, formatters
from globus_cli.types import DATA_CONTAINER_T

Expand All @@ -12,10 +12,9 @@ class LazyCurrentIdentity:
def __init__(self, value: str | None) -> None:
self._value = value

@property
def value(self) -> str:
def resolve(self, login_manager: LoginManager) -> str:
if self._value is None:
self._value = get_current_identity_id()
self._value = login_manager.get_current_identity_id()
return str(self._value)


Expand Down
7 changes: 5 additions & 2 deletions src/globus_cli/commands/collection/create/guest.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,10 @@ def collection_create_guest(

if not user_credential_id:
user_credential_id = _select_user_credential_id(
gcs_client, mapped_collection_id, local_username, identity_id.value
gcs_client,
mapped_collection_id,
local_username,
identity_id.resolve(login_manager),
)

converted_kwargs: dict[str, t.Any] = ExplicitNullType.nullify_dict(
Expand All @@ -105,7 +108,7 @@ def collection_create_guest(
"display_name": display_name,
"enable_https": enable_https,
"force_encryption": force_encryption,
"identity_id": identity_id.value,
"identity_id": identity_id.resolve(login_manager),
"info_link": info_link,
"keywords": keywords,
"mapped_collection_id": mapped_collection_id,
Expand Down
2 changes: 1 addition & 1 deletion src/globus_cli/commands/collection/create/mapped.py
Original file line number Diff line number Diff line change
Expand Up @@ -284,7 +284,7 @@ def collection_create_mapped(
"domain_name": domain_name,
"enable_https": enable_https,
"force_encryption": force_encryption,
"identity_id": identity_id.value,
"identity_id": identity_id.resolve(login_manager),
"info_link": info_link,
"keywords": keywords,
"organization": organization,
Expand Down
3 changes: 1 addition & 2 deletions src/globus_cli/commands/flows/run/resume.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
import globus_sdk

from globus_cli.login_manager import LoginManager
from globus_cli.login_manager.utils import get_current_identity_id
from globus_cli.parsing import command, run_id_arg
from globus_cli.termio import Field, display, formatters
from globus_cli.utils import CLIAuthRequirementsError
Expand Down Expand Up @@ -124,6 +123,6 @@ def _has_required_consent(
login_manager: LoginManager, required_scopes: list[str]
) -> bool:
auth_client = login_manager.get_auth_client()
user_identity_id = get_current_identity_id()
user_identity_id = login_manager.get_current_identity_id()
consents = auth_client.get_consents(user_identity_id).to_forest()
return consents.meets_scope_requirements(required_scopes)
20 changes: 6 additions & 14 deletions src/globus_cli/commands/logout.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,7 @@
import click
import globus_sdk

from globus_cli.login_manager import (
LoginManager,
delete_templated_client,
internal_native_client,
is_client_login,
remove_well_known_config,
token_storage_adapter,
)
from globus_cli.login_manager import LoginManager, is_client_login
from globus_cli.parsing import command


Expand Down Expand Up @@ -105,7 +98,7 @@ def logout_command(
# Always skip for client logins, which don't use a templated client
if delete_client and not is_client_login():
try:
delete_templated_client()
login_manager.storage.delete_templated_client()
except globus_sdk.AuthAPIError:
if not ignore_errors:
warnecho(
Expand All @@ -121,10 +114,9 @@ def logout_command(

# Attempt to revoke all tokens in storage; use the internal native client to ensure
# we have a valid Auth client
native_client = internal_native_client()
adapter = token_storage_adapter()
native_client = login_manager.storage.cli_native_client

for rs, tokendata in adapter.get_by_resource_server().items():
for rs, tokendata in login_manager.storage.adapter.get_by_resource_server().items():
for tok_key in ("access_token", "refresh_token"):
token = tokendata[tok_key]

Expand All @@ -145,9 +137,9 @@ def logout_command(
"Continuing... (--ignore-errors)",
)

adapter.remove_tokens_for_resource_server(rs)
login_manager.storage.adapter.remove_tokens_for_resource_server(rs)

remove_well_known_config("auth_user_data")
login_manager.storage.remove_well_known_config("auth_user_data")

if is_client_login():
click.echo(_CLIENT_LOGOUT_EPILOG)
Expand Down
13 changes: 3 additions & 10 deletions src/globus_cli/commands/session/show.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,7 @@

import globus_sdk

from globus_cli.login_manager import (
LoginManager,
get_client_login,
internal_auth_client,
is_client_login,
token_storage_adapter,
)
from globus_cli.login_manager import LoginManager, get_client_login, is_client_login
from globus_cli.parsing import command
from globus_cli.termio import Field, display, print_command_hint

Expand Down Expand Up @@ -54,7 +48,6 @@ def session_show(login_manager: LoginManager) -> None:
the time the user authenticated with that identity.
"""
auth_client = login_manager.get_auth_client()
adapter = token_storage_adapter()

# get a token to introspect, refreshing if necessary
try:
Expand All @@ -64,7 +57,7 @@ def session_show(login_manager: LoginManager) -> None:
except AttributeError: # if we have no RefreshTokenAuthorizor
pass

tokendata = adapter.get_token_data("auth.globus.org")
tokendata = login_manager.storage.adapter.get_token_data("auth.globus.org")
# if there's no token (e.g. not logged in), stub with empty data
if not tokendata:
session_info: dict[str, t.Any] = {}
Expand All @@ -73,7 +66,7 @@ def session_show(login_manager: LoginManager) -> None:
if is_client_login():
introspect_client = get_client_login()
else:
introspect_client = internal_auth_client()
introspect_client = login_manager.storage.cli_confidential_client

access_token = tokendata["access_token"]
res = introspect_client.oauth2_token_introspect(
Expand Down
8 changes: 5 additions & 3 deletions src/globus_cli/commands/timer/create/transfer.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@

from globus_cli.endpointish import Endpointish
from globus_cli.login_manager import LoginManager, is_client_login
from globus_cli.login_manager.utils import get_current_identity_id
from globus_cli.parsing import (
ENDPOINT_PLUS_OPTPATH,
TimedeltaType,
Expand Down Expand Up @@ -265,7 +264,9 @@ def transfer_command(
# If it's not a client login, we need to check
# that the user has the required scopes
if not is_client_login():
request_data_access = _derive_missing_scopes(auth_client, scopes_needed)
request_data_access = _derive_missing_scopes(
login_manager, auth_client, scopes_needed
)

if request_data_access:
scope_request_opts = " ".join(
Expand Down Expand Up @@ -349,11 +350,12 @@ def _derive_needed_scopes(


def _derive_missing_scopes(
login_manager: LoginManager,
auth_client: CustomAuthClient,
scopes_needed: dict[str, Scope],
) -> list[str]:
# read the identity ID stored from the login flow
user_identity_id = get_current_identity_id()
user_identity_id = login_manager.get_current_identity_id()

# get the user's Globus CLI consents
consents = auth_client.get_consents(user_identity_id).to_forest()
Expand Down
3 changes: 1 addition & 2 deletions src/globus_cli/commands/timer/resume.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
import globus_sdk

from globus_cli.login_manager import LoginManager
from globus_cli.login_manager.utils import get_current_identity_id
from globus_cli.parsing import command
from globus_cli.termio import display
from globus_cli.utils import CLIAuthRequirementsError
Expand Down Expand Up @@ -117,6 +116,6 @@ def _has_required_consent(
login_manager: LoginManager, required_scopes: list[str]
) -> bool:
auth_client = login_manager.get_auth_client()
user_identity_id = get_current_identity_id()
user_identity_id = login_manager.get_current_identity_id()
consents = auth_client.get_consents(user_identity_id).to_forest()
return consents.meets_scope_requirements(required_scopes)
16 changes: 0 additions & 16 deletions src/globus_cli/login_manager/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,29 +2,13 @@
from .errors import MissingLoginError
from .manager import LoginManager
from .scopes import compute_timer_scope
from .tokenstore import (
delete_templated_client,
internal_auth_client,
internal_native_client,
read_well_known_config,
remove_well_known_config,
store_well_known_config,
token_storage_adapter,
)
from .utils import is_remote_session

__all__ = [
"MissingLoginError",
"is_remote_session",
"LoginManager",
"delete_templated_client",
"internal_auth_client",
"internal_native_client",
"token_storage_adapter",
"is_client_login",
"get_client_login",
"store_well_known_config",
"read_well_known_config",
"remove_well_known_config",
"compute_timer_scope",
]
25 changes: 11 additions & 14 deletions src/globus_cli/login_manager/auth_flows.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,15 +8,11 @@
import globus_sdk
from globus_sdk.scopes import Scope

from .tokenstore import (
internal_auth_client,
read_well_known_config,
store_well_known_config,
token_storage_adapter,
)
from .storage import CLIStorage


def do_link_auth_flow(
storage: CLIStorage,
scopes: str | t.Sequence[str | Scope],
*,
session_params: dict[str, str] | None = None,
Expand All @@ -28,7 +24,7 @@ def do_link_auth_flow(
session_params = session_params or {}

# get the ConfidentialApp client object
auth_client = internal_auth_client()
auth_client = storage.cli_confidential_client

# start the Confidential App Grant flow
auth_client.oauth2_start_flow(
Expand All @@ -53,11 +49,12 @@ def do_link_auth_flow(
auth_code = click.prompt("Enter the resulting Authorization Code here").strip()

# finish auth flow
exchange_code_and_store(auth_client, auth_code)
exchange_code_and_store(storage, auth_client, auth_code)
return True


def do_local_server_auth_flow(
storage: CLIStorage,
scopes: str | t.Sequence[str | Scope],
*,
session_params: dict[str, str] | None = None,
Expand All @@ -78,7 +75,7 @@ def do_local_server_auth_flow(
redirect_uri = f"http://localhost:{port}"

# get the ConfidentialApp client object and start a flow
auth_client = internal_auth_client()
auth_client = storage.cli_confidential_client
auth_client.oauth2_start_flow(
refresh_tokens=True,
redirect_uri=redirect_uri,
Expand All @@ -103,11 +100,12 @@ def do_local_server_auth_flow(
click.get_current_context().exit(1)

# finish auth flow and return true
exchange_code_and_store(auth_client, auth_code)
exchange_code_and_store(storage, auth_client, auth_code)
return True


def exchange_code_and_store(
storage: CLIStorage,
auth_client: globus_sdk.ConfidentialAppAuthClient | globus_sdk.NativeAppAuthClient,
auth_code: str,
) -> None:
Expand All @@ -121,7 +119,6 @@ def exchange_code_and_store(
"""
import jwt.exceptions

adapter = token_storage_adapter()
tkn = auth_client.oauth2_exchange_code_for_tokens(auth_code)

# use a leeway of 300s
Expand Down Expand Up @@ -156,7 +153,7 @@ def exchange_code_and_store(
err=True,
)
raise
auth_user_data = read_well_known_config("auth_user_data")
auth_user_data = storage.read_well_known_config("auth_user_data")
if auth_user_data and sub_new != auth_user_data.get("sub"):
try:
for tokens in tkn.by_resource_server.values():
Expand All @@ -171,8 +168,8 @@ def exchange_code_and_store(
)
click.get_current_context().exit(1)
if not auth_user_data:
store_well_known_config("auth_user_data", {"sub": sub_new})
adapter.store(tkn)
storage.store_well_known_config("auth_user_data", {"sub": sub_new})
storage.store(tkn)


def _response_clock_delta(response: globus_sdk.GlobusHTTPResponse) -> float | None:
Expand Down
Loading

0 comments on commit d2a8289

Please sign in to comment.