Account for new EU region

logfire/_internal/cli.py

@@ -24,9 +24,8 @@
 
 from ..version import VERSION
 from .auth import DEFAULT_FILE, HOME_LOGFIRE, DefaultFile, is_logged_in, poll_for_token, request_device_code
-from .config import LogfireCredentials
+from .config import REGIONS, LogfireCredentials, get_base_url_from_token
 from .config_params import ParamManager
-from .constants import LOGFIRE_BASE_URL
 from .tracer import SDKTracerProvider
 from .utils import read_toml_file
 
@@ -55,6 +54,7 @@ def parse_whoami(args: argparse.Namespace) -> None:
     token = param_manager.load_param('token')
 
     if token:
+        base_url = base_url or get_base_url_from_token(token)
         credentials = LogfireCredentials.from_token(token, args._session, base_url)
         if credentials:
             credentials.print_token_summary()
@@ -197,16 +197,27 @@ def parse_auth(args: argparse.Namespace) -> None:
 
     This will authenticate your machine with Logfire and store the credentials.
     """
-    logfire_url = cast(str, args.logfire_url)
-
+    logfire_url = args.logfire_url
     if DEFAULT_FILE.is_file():
         data = cast(DefaultFile, read_toml_file(DEFAULT_FILE))
-        if is_logged_in(data, logfire_url):  # pragma: no branch
-            sys.stderr.write(f'You are already logged in. (Your credentials are stored in {DEFAULT_FILE})\n')
-            return
     else:
         data: DefaultFile = {'tokens': {}}
 
+    if logfire_url:
+        logged_in = is_logged_in(data, logfire_url)
+    else:
+        logged_in = any(is_logged_in(data, url) for url in data['tokens'])
+
+    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',
@@ -215,6 +226,20 @@ def parse_auth(args: argparse.Namespace) -> None:
             '\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
@@ -245,8 +270,7 @@ def parse_auth(args: argparse.Namespace) -> None:
 
 def parse_list_projects(args: argparse.Namespace) -> None:
     """List user projects."""
-    logfire_url = args.logfire_url
-    projects = LogfireCredentials.get_user_projects(session=args._session, logfire_api_url=logfire_url)
+    projects = LogfireCredentials.get_user_projects(session=args._session, logfire_api_url=args.logfire_url)
     if projects:
         sys.stderr.write(
             _pretty_table(
@@ -276,6 +300,8 @@ def parse_create_new_project(args: argparse.Namespace) -> None:
     """Create a new project."""
     data_dir = Path(args.data_dir)
     logfire_url = args.logfire_url
+    if logfire_url is None:  # pragma: no cover
+        _, logfire_url = LogfireCredentials._get_user_token_data()  # type: ignore
     project_name = args.project_name
     organization = args.org
     default_organization = args.default_org
@@ -294,6 +320,8 @@ def parse_use_project(args: argparse.Namespace) -> None:
     """Use an existing project."""
     data_dir = Path(args.data_dir)
     logfire_url = args.logfire_url
+    if logfire_url is None:  # pragma: no cover
+        _, logfire_url = LogfireCredentials._get_user_token_data()  # type: ignore
     project_name = args.project_name
     organization = args.org
 
@@ -374,6 +402,13 @@ def _pretty_table(header: list[str], rows: list[list[str]]):
     return '\n'.join(lines) + '\n'
 
 
+def _get_logfire_url(logfire_url: str | None, region: str | None) -> str | None:
+    if logfire_url is not None:
+        return logfire_url
+    if region is not None:
+        return REGIONS[region]['base_url']
+
+
 class SplitArgs(argparse.Action):
     def __call__(
         self,
@@ -397,7 +432,9 @@ def _main(args: list[str] | None = None) -> None:
 
     parser.add_argument('--version', action='store_true', help='show the version and exit')
     global_opts = parser.add_argument_group(title='global options')
-    global_opts.add_argument('--logfire-url', default=LOGFIRE_BASE_URL, help=argparse.SUPPRESS)
+    url_or_region_grp = global_opts.add_mutually_exclusive_group()
+    url_or_region_grp.add_argument('--logfire-url', help=argparse.SUPPRESS)
+    url_or_region_grp.add_argument('--region', choices=REGIONS, help='the region to use')
     parser.set_defaults(func=lambda _: parser.print_help())  # type: ignore
     subparsers = parser.add_subparsers(title='commands', metavar='')
 
@@ -444,6 +481,7 @@ def _main(args: list[str] | None = None) -> None:
     cmd_info.set_defaults(func=parse_info)
 
     namespace = parser.parse_args(args)
+    namespace.logfire_url = _get_logfire_url(namespace.logfire_url, namespace.region)
 
     trace.set_tracer_provider(tracer_provider=SDKTracerProvider())
     tracer = trace.get_tracer(__name__)