From 16c1191e6539e66a134e84ad4f9c102c0e45f544 Mon Sep 17 00:00:00 2001 From: "Michiel De Smet (aider)" Date: Thu, 12 Jun 2025 14:38:46 +0800 Subject: [PATCH 1/6] refactor: add auth_options decorator for CLI commands --- src/datapilot/cli/decorators.py | 90 +++++++++++++++++++++ src/datapilot/cli/main.py | 80 +----------------- src/datapilot/core/knowledge/cli.py | 11 ++- src/datapilot/core/platforms/dbt/cli/cli.py | 24 +++--- 4 files changed, 111 insertions(+), 94 deletions(-) create mode 100644 src/datapilot/cli/decorators.py diff --git a/src/datapilot/cli/decorators.py b/src/datapilot/cli/decorators.py new file mode 100644 index 0000000..bd2af18 --- /dev/null +++ b/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 diff --git a/src/datapilot/cli/main.py b/src/datapilot/cli/main.py index 87a6135..a2aee93 100644 --- a/src/datapilot/cli/main.py +++ b/src/datapilot/cli/main.py @@ -1,10 +1,4 @@ -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 @@ -12,84 +6,14 @@ 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) diff --git a/src/datapilot/core/knowledge/cli.py b/src/datapilot/core/knowledge/cli.py index 2fe30f8..fa4c28d 100644 --- a/src/datapilot/core/knowledge/cli.py +++ b/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( diff --git a/src/datapilot/core/platforms/dbt/cli/cli.py b/src/datapilot/core/platforms/dbt/cli/cli.py index 0f63167..b60f18c 100644 --- a/src/datapilot/core/platforms/dbt/cli/cli.py +++ b/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: From 73bbbb31b3006bd7b67454e47c7c630fa1233537 Mon Sep 17 00:00:00 2001 From: Michiel De Smet Date: Thu, 12 Jun 2025 14:42:00 +0800 Subject: [PATCH 2/6] style: format auth_options decorator --- src/datapilot/cli/decorators.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/datapilot/cli/decorators.py b/src/datapilot/cli/decorators.py index bd2af18..793ff82 100644 --- a/src/datapilot/cli/decorators.py +++ b/src/datapilot/cli/decorators.py @@ -49,6 +49,7 @@ def process_config(config): 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") @@ -86,5 +87,5 @@ def wrapper(ctx, token, instance_name, backend_url, *args, **kwargs): ctx.obj.setdefault("backend_url", "https://api.myaltimate.com") return f(ctx, *args, **kwargs) - + return wrapper From 6b4334b3cc5ac79425b0f270b00a2bc82e9d2cf6 Mon Sep 17 00:00:00 2001 From: "Michiel De Smet (aider)" Date: Thu, 12 Jun 2025 14:52:28 +0800 Subject: [PATCH 3/6] refactor: remove click context usage in favor of direct parameter passing --- src/datapilot/cli/decorators.py | 40 ++++++++++----------- src/datapilot/cli/main.py | 5 +-- src/datapilot/core/knowledge/cli.py | 10 ++---- src/datapilot/core/platforms/dbt/cli/cli.py | 18 ++++------ 4 files changed, 29 insertions(+), 44 deletions(-) diff --git a/src/datapilot/cli/decorators.py b/src/datapilot/cli/decorators.py index 793ff82..ed8d366 100644 --- a/src/datapilot/cli/decorators.py +++ b/src/datapilot/cli/decorators.py @@ -54,7 +54,7 @@ def auth_options(f): @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): + def wrapper(token, instance_name, backend_url, *args, **kwargs): # Load .env file from current directory if it exists load_dotenv() @@ -65,27 +65,27 @@ def wrapper(ctx, token, instance_name, backend_url, *args, **kwargs): # 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, then override with CLI arguments if provided + final_token = token + final_instance_name = instance_name + final_backend_url = backend_url - # 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 + # Use file config if CLI argument not provided + if final_token is None and "altimateApiKey" in file_config: + final_token = file_config["altimateApiKey"] + if final_instance_name is None and "altimateInstanceName" in file_config: + final_instance_name = file_config["altimateInstanceName"] + if final_backend_url == "https://api.myaltimate.com" and "altimateUrl" in file_config: + final_backend_url = file_config["altimateUrl"] # 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) + if final_token is None: + final_token = None + if final_instance_name is None: + final_instance_name = None + if final_backend_url is None: + final_backend_url = "https://api.myaltimate.com" + + return f(final_token, final_instance_name, final_backend_url, *args, **kwargs) return wrapper diff --git a/src/datapilot/cli/main.py b/src/datapilot/cli/main.py index a2aee93..dcf0a92 100644 --- a/src/datapilot/cli/main.py +++ b/src/datapilot/cli/main.py @@ -8,11 +8,8 @@ @click.group() @click.version_option(version=__version__, prog_name="datapilot") -@click.pass_context -def datapilot(ctx): +def datapilot(): """Altimate CLI for DBT project management.""" - # Ensure context object exists - ctx.ensure_object(dict) datapilot.add_command(dbt) diff --git a/src/datapilot/core/knowledge/cli.py b/src/datapilot/core/knowledge/cli.py index fa4c28d..bdb5395 100644 --- a/src/datapilot/core/knowledge/cli.py +++ b/src/datapilot/core/knowledge/cli.py @@ -15,19 +15,13 @@ 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): +def serve(token, instance_name, backend_url, port): """Serve knowledge bases via HTTP server.""" - # 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( "Error: API token and instance name are required. Use --token and --instance-name options or set them in config.", err=True ) - ctx.exit(1) + raise click.Abort() # Set context data for the handler KnowledgeBaseHandler.token = token diff --git a/src/datapilot/core/platforms/dbt/cli/cli.py b/src/datapilot/core/platforms/dbt/cli/cli.py index b60f18c..c7289c9 100644 --- a/src/datapilot/core/platforms/dbt/cli/cli.py +++ b/src/datapilot/core/platforms/dbt/cli/cli.py @@ -57,9 +57,10 @@ def dbt(): default=None, help="Selective model testing. Specify one or more models to run tests on.", ) -@click.pass_context def project_health( - ctx, + token, + instance_name, + backend_url, manifest_path, catalog_path, config_path=None, @@ -70,10 +71,6 @@ def project_health( Validate the DBT project's configuration and structure. :param manifest_path: Path to the DBT manifest file. """ - # 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: @@ -141,19 +138,16 @@ def project_health( ) @click.option("--manifest-path", required=True, prompt="Manifest Path", help="Path to the manifest file.") @click.option("--catalog-path", required=False, prompt=False, help="Path to the catalog file.") -@click.pass_context def onboard( - ctx, + token, + instance_name, + backend_url, dbt_core_integration_id, dbt_core_integration_environment, manifest_path, catalog_path, ): """Onboard a manifest file to DBT.""" - # 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: From 34a829ca7582b65ea810be28e8205dd8a23094f5 Mon Sep 17 00:00:00 2001 From: Michiel De Smet Date: Thu, 12 Jun 2025 14:56:10 +0800 Subject: [PATCH 4/6] refactor: remove unused config mapping in auth decorator --- src/datapilot/cli/decorators.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/datapilot/cli/decorators.py b/src/datapilot/cli/decorators.py index ed8d366..6267f4c 100644 --- a/src/datapilot/cli/decorators.py +++ b/src/datapilot/cli/decorators.py @@ -62,9 +62,6 @@ def wrapper(token, instance_name, backend_url, *args, **kwargs): 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"} - # Apply file config first, then override with CLI arguments if provided final_token = token final_instance_name = instance_name From ee8d1a5b70d0b2afa314a3e1e405aa830ebcc577 Mon Sep 17 00:00:00 2001 From: Michiel De Smet Date: Thu, 12 Jun 2025 15:09:18 +0800 Subject: [PATCH 5/6] fix: return instead of raising click.Abort in serve --- src/datapilot/core/knowledge/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/datapilot/core/knowledge/cli.py b/src/datapilot/core/knowledge/cli.py index bdb5395..bf43c5a 100644 --- a/src/datapilot/core/knowledge/cli.py +++ b/src/datapilot/core/knowledge/cli.py @@ -21,7 +21,7 @@ def serve(token, instance_name, backend_url, port): click.echo( "Error: API token and instance name are required. Use --token and --instance-name options or set them in config.", err=True ) - raise click.Abort() + return # Set context data for the handler KnowledgeBaseHandler.token = token From ace15212e91d8dd3e6353091dd779eeed1e5b645 Mon Sep 17 00:00:00 2001 From: Michiel De Smet Date: Thu, 12 Jun 2025 15:11:58 +0800 Subject: [PATCH 6/6] fix: raise click.Abort on missing token and instance name --- src/datapilot/core/knowledge/cli.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/datapilot/core/knowledge/cli.py b/src/datapilot/core/knowledge/cli.py index bf43c5a..3119a21 100644 --- a/src/datapilot/core/knowledge/cli.py +++ b/src/datapilot/core/knowledge/cli.py @@ -21,7 +21,7 @@ def serve(token, instance_name, backend_url, port): click.echo( "Error: API token and instance name are required. Use --token and --instance-name options or set them in config.", err=True ) - return + raise click.Abort # Set context data for the handler KnowledgeBaseHandler.token = token