#!/usr/bin/env python3 """ Generic Snowflake query executor via CLI. For Cursor AI to execute queries. This script uses the DBT_DEV_WH environment variable if available to specify the Snowflake warehouse. Otherwise, it defaults to Warehouse.SMALL. Usage: # Execute a query directly python query_snowflake.py --query "SELECT * FROM suno_prod.prod.clip LIMIT 10" # Execute a query from a file python query_snowflake.py --file query.sql # Query a specific table with auto-generated SELECT python query_snowflake.py --table clip --database suno_prod --schema prod --limit 10 # Query with custom output format python query_snowflake.py --query "SELECT * FROM table" --output csv python query_snowflake.py --query "SELECT * FROM table" --output json """ import sys import os import argparse from pathlib import Path from typing import Optional from dotenv import load_dotenv # Load environment variables from .env file env_path = Path(__file__).parent / ".env" if env_path.exists(): load_dotenv(env_path) # Handle DBT_SNOWFLAKE_PRIVATE_KEY_PATH if set # This allows using dbt dev environment configuration if os.environ.get("DBT_SNOWFLAKE_PRIVATE_KEY_PATH") and not os.environ.get("SNOWFLAKE_PRIVATE_KEY"): private_key_path = os.environ.get("DBT_SNOWFLAKE_PRIVATE_KEY_PATH") try: with open(private_key_path, "r") as f: os.environ["SNOWFLAKE_PRIVATE_KEY"] = f.read() except Exception as e: print(f"Warning: Could not read private key from {private_key_path}: {e}", file=sys.stderr) # Set DBT_USER as SNOWFLAKE_ACCOUNT_USER if not already set if os.environ.get("DBT_USER") and not os.environ.get("SNOWFLAKE_ACCOUNT_USER"): os.environ["SNOWFLAKE_ACCOUNT_USER"] = os.environ["DBT_USER"] # Add src to path sys.path.insert(0, str(Path(__file__).parent / "src")) from src.utils.database import query_snowflake, get_snowflake_private_key def execute_query( query: str, output_format: str = "table", max_rows: Optional[int] = None, ) -> None: """Execute a Snowflake query and display results.""" print(f"\n{'='*80}") print(f"Executing query:") print(f"{'='*80}") print(query) print(f"{'='*80}\n") try: # Use DBT_DEV_WH if available, otherwise use default Warehouse.SMALL if os.environ.get("DBT_DEV_WH"): # If DBT_DEV_WH is set, create a custom connection with the specified warehouse import pandas as pd import snowflake.connector # Get warehouse name from environment warehouse_name = os.environ["DBT_DEV_WH"] # Create connection with custom warehouse using existing utility functions connection_params = { "user": os.environ["SNOWFLAKE_ACCOUNT_USER"], "private_key": get_snowflake_private_key(), "account": os.environ["SNOWFLAKE_ACCOUNT"], "warehouse": warehouse_name, "database": "SUNO_PROD", "schema": "PROD", } # Use the connection to execute query with snowflake.connector.connect(**connection_params) as conn: cursor = conn.cursor() cursor.execute(query) results = cursor.fetchall() columns = [desc[0] for desc in cursor.description] df = pd.DataFrame(results, columns=columns) else: # Use default query_snowflake function df = query_snowflake(query) if len(df) == 0: print("Query returned 0 rows.") return print(f"Rows returned: {len(df)}") print(f"Columns: {list(df.columns)}\n") # Limit rows for display if specified display_df = df.head(max_rows) if max_rows else df if output_format == "csv": print(display_df.to_csv(index=False)) elif output_format == "json": print(display_df.to_json(orient="records", indent=2)) elif output_format == "markdown": try: print(display_df.to_markdown(index=False)) except AttributeError: # Fallback for older pandas versions without to_markdown print(display_df.to_string(index=False)) else: # table (default) print(display_df.to_string(index=False)) if max_rows and len(df) > max_rows: print(f"\n... ({len(df) - max_rows} more rows not shown)") except Exception as e: print(f"Error executing query: {e}", file=sys.stderr) sys.exit(1) def read_query_file(file_path: str) -> str: """Read SQL query from a file.""" path = Path(file_path) if not path.exists(): raise FileNotFoundError(f"Query file not found: {file_path}") return path.read_text() def main(): parser = argparse.ArgumentParser( description="Execute Snowflake queries via CLI", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=__doc__ ) # Query input options (mutually exclusive) query_group = parser.add_mutually_exclusive_group(required=True) query_group.add_argument( "--query", "-q", type=str, help="SQL query to execute directly" ) query_group.add_argument( "--file", "-f", type=str, help="Path to SQL file containing query" ) query_group.add_argument( "--table", "-t", type=str, help="Table name (auto-generates SELECT * FROM query)" ) # Table query options (only used with --table) parser.add_argument( "--database", "-d", type=str, default="suno_prod", help="Database name (default: suno_prod)" ) parser.add_argument( "--schema", "-s", type=str, default="prod", help="Schema name (default: prod)" ) parser.add_argument( "--limit", "-l", type=int, default=10, help="Limit for table query (default: 10)" ) # Output options parser.add_argument( "--output", "-o", type=str, choices=["table", "csv", "json", "markdown"], default="table", help="Output format (default: table)" ) parser.add_argument( "--max-rows", type=int, help="Maximum rows to display (default: all)" ) args = parser.parse_args() # Determine the query to execute if args.query: query = args.query elif args.file: query = read_query_file(args.file) elif args.table: query = f""" SELECT * FROM {args.database}.{args.schema}.{args.table} LIMIT {args.limit} """ else: parser.error("Must provide --query, --file, or --table") # Execute the query execute_query(query, output_format=args.output, max_rows=args.max_rows) if __name__ == "__main__": main()