refactor: add auth_options decorator for CLI commands

Author: mdesmet

src/datapilot/cli/decorators.py

@@ -0,0 +1,90 @@
+import json
+import os
+import re
+from functools import wraps
+from pathlib import Path
+
+import click
+from dotenv import load_dotenv
+
+
+def load_config_from_file():
+    """Load configuration from ~/.altimate/altimate.json if it exists."""
+    config_path = Path.home() / ".altimate" / "altimate.json"
+
+    if not config_path.exists():
+        return {}
+
+    try:
+        with config_path.open() as f:
+            config = json.load(f)
+        return config
+    except (OSError, json.JSONDecodeError) as e:
+        click.echo(f"Warning: Failed to load config from {config_path}: {e}", err=True)
+        return {}
+
+
+def substitute_env_vars(value):
+    """Replace ${env:ENV_VARIABLE} patterns with actual environment variable values."""
+    if not isinstance(value, str):
+        return value
+
+    # Pattern to match ${env:VARIABLE_NAME}
+    pattern = r"\$\{env:([^}]+)\}"
+
+    def replacer(match):
+        env_var = match.group(1)
+        return os.environ.get(env_var, match.group(0))
+
+    return re.sub(pattern, replacer, value)
+
+
+def process_config(config):
+    """Process configuration dictionary to substitute environment variables."""
+    processed = {}
+    for key, value in config.items():
+        processed[key] = substitute_env_vars(value)
+    return processed
+
+
+def auth_options(f):
+    """Decorator to add authentication options to commands."""
+    @click.option("--token", required=False, help="Your API token for authentication.", hide_input=True)
+    @click.option("--instance-name", required=False, help="Your tenant ID.")
+    @click.option("--backend-url", required=False, help="Altimate's Backend URL", default="https://api.myaltimate.com")
+    @wraps(f)
+    def wrapper(ctx, token, instance_name, backend_url, *args, **kwargs):
+        # Load .env file from current directory if it exists
+        load_dotenv()
+
+        # Load configuration from file
+        file_config = load_config_from_file()
+        file_config = process_config(file_config)
+
+        # Map config file keys to CLI option names
+        config_mapping = {"altimateApiKey": "token", "altimateInstanceName": "instance_name", "altimateUrl": "backend_url"}
+
+        # Store common options in context
+        ctx.ensure_object(dict)
+
+        # Apply file config first
+        for file_key, cli_key in config_mapping.items():
+            if file_key in file_config:
+                ctx.obj[cli_key] = file_config[file_key]
+
+        # Override with CLI arguments if provided
+        if token is not None:
+            ctx.obj["token"] = token
+        if instance_name is not None:
+            ctx.obj["instance_name"] = instance_name
+        if backend_url != "https://api.myaltimate.com":  # Only override if not default
+            ctx.obj["backend_url"] = backend_url
+
+        # Set defaults if nothing was provided
+        ctx.obj.setdefault("token", None)
+        ctx.obj.setdefault("instance_name", None)
+        ctx.obj.setdefault("backend_url", "https://api.myaltimate.com")
+
+        return f(ctx, *args, **kwargs)
+    
+    return wrapper

src/datapilot/cli/main.py

@@ -1,95 +1,19 @@
-import json
-import os
-import re
-from pathlib import Path
-
 import click
-from dotenv import load_dotenv
 
 from datapilot import __version__
 from datapilot.core.knowledge.cli import cli as knowledge
 from datapilot.core.mcp_utils.mcp import mcp
 from datapilot.core.platforms.dbt.cli.cli import dbt
 
 
-def load_config_from_file():
-    """Load configuration from ~/.altimate/altimate.json if it exists."""
-    config_path = Path.home() / ".altimate" / "altimate.json"
-
-    if not config_path.exists():
-        return {}
-
-    try:
-        with config_path.open() as f:
-            config = json.load(f)
-        return config
-    except (OSError, json.JSONDecodeError) as e:
-        click.echo(f"Warning: Failed to load config from {config_path}: {e}", err=True)
-        return {}
-
-
-def substitute_env_vars(value):
-    """Replace ${env:ENV_VARIABLE} patterns with actual environment variable values."""
-    if not isinstance(value, str):
-        return value
-
-    # Pattern to match ${env:VARIABLE_NAME}
-    pattern = r"\$\{env:([^}]+)\}"
-
-    def replacer(match):
-        env_var = match.group(1)
-        return os.environ.get(env_var, match.group(0))
-
-    return re.sub(pattern, replacer, value)
-
-
-def process_config(config):
-    """Process configuration dictionary to substitute environment variables."""
-    processed = {}
-    for key, value in config.items():
-        processed[key] = substitute_env_vars(value)
-    return processed
-
-
 @click.group()
 @click.version_option(version=__version__, prog_name="datapilot")
-@click.option("--token", required=False, help="Your API token for authentication.", hide_input=True)
-@click.option("--instance-name", required=False, help="Your tenant ID.")
-@click.option("--backend-url", required=False, help="Altimate's Backend URL", default="https://api.myaltimate.com")
 @click.pass_context
-def datapilot(ctx, token, instance_name, backend_url):
+def datapilot(ctx):
     """Altimate CLI for DBT project management."""
-    # Load .env file from current directory if it exists
-    load_dotenv()
-
-    # Load configuration from file
-    file_config = load_config_from_file()
-    file_config = process_config(file_config)
-
-    # Map config file keys to CLI option names
-    config_mapping = {"altimateApiKey": "token", "altimateInstanceName": "instance_name", "altimateUrl": "backend_url"}
-
-    # Store common options in context, with CLI args taking precedence
+    # Ensure context object exists
     ctx.ensure_object(dict)
 
-    # Apply file config first
-    for file_key, cli_key in config_mapping.items():
-        if file_key in file_config:
-            ctx.obj[cli_key] = file_config[file_key]
-
-    # Override with CLI arguments if provided
-    if token is not None:
-        ctx.obj["token"] = token
-    if instance_name is not None:
-        ctx.obj["instance_name"] = instance_name
-    if backend_url != "https://api.myaltimate.com":  # Only override if not default
-        ctx.obj["backend_url"] = backend_url
-
-    # Set defaults if nothing was provided
-    ctx.obj.setdefault("token", None)
-    ctx.obj.setdefault("instance_name", None)
-    ctx.obj.setdefault("backend_url", "https://api.myaltimate.com")
-
 
 datapilot.add_command(dbt)
 datapilot.add_command(mcp)

src/datapilot/core/knowledge/cli.py

@@ -2,6 +2,8 @@
 
 import click
 
+from datapilot.cli.decorators import auth_options
+
 from .server import KnowledgeBaseHandler
 
 
@@ -11,14 +13,15 @@ def cli():
 
 
 @cli.command()
+@auth_options
 @click.option("--port", default=4000, help="Port to run the server on")
 @click.pass_context
 def serve(ctx, port):
     """Serve knowledge bases via HTTP server."""
-    # Get configuration from parent context
-    token = ctx.parent.obj.get("token")
-    instance_name = ctx.parent.obj.get("instance_name")
-    backend_url = ctx.parent.obj.get("backend_url")
+    # Get configuration from context
+    token = ctx.obj.get("token")
+    instance_name = ctx.obj.get("instance_name")
+    backend_url = ctx.obj.get("backend_url")
 
     if not token or not instance_name:
         click.echo(

src/datapilot/core/platforms/dbt/cli/cli.py

@@ -2,6 +2,7 @@
 
 import click
 
+from datapilot.cli.decorators import auth_options
 from datapilot.clients.altimate.utils import check_token_and_instance
 from datapilot.clients.altimate.utils import get_all_dbt_configs
 from datapilot.clients.altimate.utils import onboard_file
@@ -24,14 +25,12 @@
 
 # New dbt group
 @click.group()
-@click.pass_context
-def dbt(ctx):
+def dbt():
     """DBT specific commands."""
-    # Ensure context object exists
-    ctx.ensure_object(dict)
 
 
 @dbt.command("project-health")
+@auth_options
 @click.option(
     "--manifest-path",
     required=True,
@@ -71,10 +70,10 @@ def project_health(
     Validate the DBT project's configuration and structure.
     :param manifest_path: Path to the DBT manifest file.
     """
-    # Get common options from parent context
-    token = ctx.parent.obj.get("token")
-    instance_name = ctx.parent.obj.get("instance_name")
-    backend_url = ctx.parent.obj.get("backend_url")
+    # Get common options from context
+    token = ctx.obj.get("token")
+    instance_name = ctx.obj.get("instance_name")
+    backend_url = ctx.obj.get("backend_url")
 
     config = None
     if config_path:
@@ -135,6 +134,7 @@ def project_health(
 
 
 @dbt.command("onboard")
+@auth_options
 @click.option("--dbt_core_integration_id", prompt="DBT Core Integration ID", help="DBT Core Integration ID")
 @click.option(
     "--dbt_core_integration_environment", default="PROD", prompt="DBT Core Integration Environment", help="DBT Core Integration Environment"
@@ -150,10 +150,10 @@ def onboard(
     catalog_path,
 ):
     """Onboard a manifest file to DBT."""
-    # Get common options from parent context
-    token = ctx.parent.obj.get("token")
-    instance_name = ctx.parent.obj.get("instance_name")
-    backend_url = ctx.parent.obj.get("backend_url")
+    # Get common options from context
+    token = ctx.obj.get("token")
+    instance_name = ctx.obj.get("instance_name")
+    backend_url = ctx.obj.get("backend_url")
 
     # For onboard command, token and instance_name are required
     if not token: