Trigger auth command from prompt (#1423)

Author: Kludex

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 <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)

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 <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')

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)

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)