diff --git a/logfire/_internal/cli/__init__.py b/logfire/_internal/cli/__init__.py index a2c2333ef..042905f88 100644 --- a/logfire/_internal/cli/__init__.py +++ b/logfire/_internal/cli/__init__.py @@ -8,27 +8,26 @@ import platform import sys import warnings -import webbrowser from collections.abc import Sequence from operator import itemgetter from pathlib import Path from typing import Any -from urllib.parse import urlparse import requests from opentelemetry import trace from rich.console import Console -from logfire._internal.cli.prompt import parse_prompt from logfire.exceptions import LogfireConfigError from logfire.propagate import ContextCarrier, get_context from ...version import VERSION -from ..auth import DEFAULT_FILE, HOME_LOGFIRE, UserTokenCollection, poll_for_token, request_device_code +from ..auth import HOME_LOGFIRE from ..client import LogfireClient from ..config import REGIONS, LogfireCredentials, get_base_url_from_token from ..config_params import ParamManager from ..tracer import SDKTracerProvider +from .auth import parse_auth +from .prompt import parse_prompt from .run import collect_instrumentation_context, parse_run, print_otel_summary BASE_OTEL_INTEGRATION_URL = 'https://opentelemetry-python-contrib.readthedocs.io/en/latest/instrumentation/' @@ -117,68 +116,6 @@ def parse_inspect(args: argparse.Namespace) -> None: console.print('No recommended packages found. You are all set!', style='green') # pragma: no cover -def parse_auth(args: argparse.Namespace) -> None: - """Authenticate with Logfire. - - This will authenticate your machine with Logfire and store the credentials. - """ - logfire_url: str | None = args.logfire_url - - tokens_collection = UserTokenCollection() - logged_in = tokens_collection.is_logged_in(logfire_url) - - if logged_in: - sys.stderr.writelines( - ( - f'You are already logged in. (Your credentials are stored in {DEFAULT_FILE})\n', - 'If you would like to log in using a different account, use the --region argument:\n', - 'logfire --region auth\n', - ) - ) - return - - sys.stderr.writelines( - ( - '\n', - 'Welcome to Logfire! 🔥\n', - 'Before you can send data to Logfire, we need to authenticate you.\n', - '\n', - ) - ) - if not logfire_url: - selected_region = -1 - while not (1 <= selected_region <= len(REGIONS)): - sys.stderr.write('Logfire is available in multiple data regions. Please select one:\n') - for i, (region_id, region_data) in enumerate(REGIONS.items(), start=1): - sys.stderr.write(f'{i}. {region_id.upper()} (GCP region: {region_data["gcp_region"]})\n') - - try: - selected_region = int( - input(f'Selected region [{"/".join(str(i) for i in range(1, len(REGIONS) + 1))}]: ') - ) - except ValueError: - selected_region = -1 - logfire_url = list(REGIONS.values())[selected_region - 1]['base_url'] - - device_code, frontend_auth_url = request_device_code(args._session, logfire_url) - frontend_host = urlparse(frontend_auth_url).netloc - input(f'Press Enter to open {frontend_host} in your browser...') - try: - webbrowser.open(frontend_auth_url, new=2) - except webbrowser.Error: - pass - sys.stderr.writelines( - ( - f"Please open {frontend_auth_url} in your browser to authenticate if it hasn't already.\n", - 'Waiting for you to authenticate with Logfire...\n', - ) - ) - - tokens_collection.add_token(logfire_url, poll_for_token(args._session, device_code, logfire_url)) - sys.stderr.write('Successfully authenticated!\n') - sys.stderr.write(f'\nYour Logfire credentials are stored in {DEFAULT_FILE}\n') - - def parse_list_projects(args: argparse.Namespace) -> None: """List user projects.""" client = LogfireClient.from_url(args.logfire_url) diff --git a/logfire/_internal/cli/auth.py b/logfire/_internal/cli/auth.py new file mode 100644 index 000000000..ddb7529bf --- /dev/null +++ b/logfire/_internal/cli/auth.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +import argparse +import sys +import webbrowser +from urllib.parse import urlparse + +from ..auth import DEFAULT_FILE, UserTokenCollection, poll_for_token, request_device_code +from ..config import REGIONS + + +def parse_auth(args: argparse.Namespace) -> None: + """Authenticate with Logfire. + + This will authenticate your machine with Logfire and store the credentials. + """ + logfire_url: str | None = args.logfire_url + + tokens_collection = UserTokenCollection() + logged_in = tokens_collection.is_logged_in(logfire_url) + + if logged_in: + sys.stderr.writelines( + ( + f'You are already logged in. (Your credentials are stored in {DEFAULT_FILE})\n', + 'If you would like to log in using a different account, use the --region argument:\n', + 'logfire --region auth\n', + ) + ) + return + + sys.stderr.writelines( + ( + '\n', + 'Welcome to Logfire! 🔥\n', + 'Before you can send data to Logfire, we need to authenticate you.\n', + '\n', + ) + ) + if not logfire_url: + selected_region = -1 + while not (1 <= selected_region <= len(REGIONS)): + sys.stderr.write('Logfire is available in multiple data regions. Please select one:\n') + for i, (region_id, region_data) in enumerate(REGIONS.items(), start=1): + sys.stderr.write(f'{i}. {region_id.upper()} (GCP region: {region_data["gcp_region"]})\n') + + try: + selected_region = int( + input(f'Selected region [{"/".join(str(i) for i in range(1, len(REGIONS) + 1))}]: ') + ) + except ValueError: + selected_region = -1 + logfire_url = list(REGIONS.values())[selected_region - 1]['base_url'] + + device_code, frontend_auth_url = request_device_code(args._session, logfire_url) + frontend_host = urlparse(frontend_auth_url).netloc + + # We are not using the `prompt` parameter from `input` here because we want to write to stderr. + sys.stderr.write(f'Press Enter to open {frontend_host} in your browser...\n') + input() + + try: + webbrowser.open(frontend_auth_url, new=2) + except webbrowser.Error: + pass + sys.stderr.writelines( + ( + f"Please open {frontend_auth_url} in your browser to authenticate if it hasn't already.\n", + 'Waiting for you to authenticate with Logfire...\n', + ) + ) + + tokens_collection.add_token(logfire_url, poll_for_token(args._session, device_code, logfire_url)) + sys.stderr.write('Successfully authenticated!\n') + sys.stderr.write(f'\nYour Logfire credentials are stored in {DEFAULT_FILE}\n') diff --git a/logfire/_internal/cli/prompt.py b/logfire/_internal/cli/prompt.py index fe21f0589..5b39c8687 100644 --- a/logfire/_internal/cli/prompt.py +++ b/logfire/_internal/cli/prompt.py @@ -14,6 +14,7 @@ from rich.console import Console +from logfire._internal.cli.auth import parse_auth from logfire._internal.client import LogfireClient from logfire.exceptions import LogfireConfigError @@ -34,9 +35,9 @@ def parse_prompt(args: argparse.Namespace) -> None: try: client = LogfireClient.from_url(args.logfire_url) - except LogfireConfigError as e: # pragma: no cover - console.print(e.args[0], style='red') - return + except LogfireConfigError: # pragma: no cover + parse_auth(args) + client = LogfireClient.from_url(args.logfire_url) if args.claude: configure_claude(client, args.organization, args.project, console) diff --git a/tests/test_cli.py b/tests/test_cli.py index 49452518f..7199e0381 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -333,8 +333,8 @@ def test_auth(tmp_path: Path, webbrowser_error: bool, capsys: pytest.CaptureFixt with ExitStack() as stack: stack.enter_context(patch('logfire._internal.auth.DEFAULT_FILE', auth_file)) # Necessary to assert that credentials are written to the `auth_file` (which happens from the `cli` module) - stack.enter_context(patch('logfire._internal.cli.DEFAULT_FILE', auth_file)) - stack.enter_context(patch('logfire._internal.cli.input')) + stack.enter_context(patch('logfire._internal.cli.auth.DEFAULT_FILE', auth_file)) + stack.enter_context(patch('logfire._internal.cli.auth.input')) webbrowser_open = stack.enter_context( patch('webbrowser.open', side_effect=webbrowser.Error if webbrowser_error is True else None) ) @@ -369,6 +369,7 @@ def test_auth(tmp_path: Path, webbrowser_error: bool, capsys: pytest.CaptureFixt 'Welcome to Logfire! 🔥', 'Before you can send data to Logfire, we need to authenticate you.', '', + 'Press Enter to open example.com in your browser...', "Please open http://example.com/auth in your browser to authenticate if it hasn't already.", 'Waiting for you to authenticate with Logfire...', 'Successfully authenticated!', @@ -384,8 +385,8 @@ def test_auth_temp_failure(tmp_path: Path) -> None: auth_file = tmp_path / 'default.toml' with ExitStack() as stack: stack.enter_context(patch('logfire._internal.auth.DEFAULT_FILE', auth_file)) - stack.enter_context(patch('logfire._internal.cli.input')) - stack.enter_context(patch('logfire._internal.cli.webbrowser.open')) + stack.enter_context(patch('logfire._internal.cli.auth.input')) + stack.enter_context(patch('logfire._internal.cli.auth.webbrowser.open')) m = requests_mock.Mocker() stack.enter_context(m) @@ -409,8 +410,8 @@ def test_auth_permanent_failure(tmp_path: Path) -> None: auth_file = tmp_path / 'default.toml' with ExitStack() as stack: stack.enter_context(patch('logfire._internal.auth.DEFAULT_FILE', auth_file)) - stack.enter_context(patch('logfire._internal.cli.input')) - stack.enter_context(patch('logfire._internal.cli.webbrowser.open')) + stack.enter_context(patch('logfire._internal.cli.auth.input')) + stack.enter_context(patch('logfire._internal.cli.auth.webbrowser.open')) m = requests_mock.Mocker() stack.enter_context(m) @@ -439,11 +440,11 @@ def test_auth_no_region_specified(tmp_path: Path) -> None: with ExitStack() as stack: stack.enter_context(patch('logfire._internal.auth.DEFAULT_FILE', auth_file)) # Necessary to assert that credentials are written to the `auth_file` (which happens from the `cli` module) - stack.enter_context(patch('logfire._internal.cli.DEFAULT_FILE', auth_file)) + stack.enter_context(patch('logfire._internal.cli.auth.DEFAULT_FILE', auth_file)) # 'not_an_int' is used as the first input to test that invalid inputs are supported, # '2' will result in the EU region being used: - stack.enter_context(patch('logfire._internal.cli.input', side_effect=['not_an_int', '2', ''])) - stack.enter_context(patch('logfire._internal.cli.webbrowser.open')) + stack.enter_context(patch('logfire._internal.cli.auth.input', side_effect=['not_an_int', '2', ''])) + stack.enter_context(patch('logfire._internal.cli.auth.webbrowser.open')) m = requests_mock.Mocker() stack.enter_context(m)