ds-agent-cli 0.1.0

package/src/tools/cloud_data_sources.py

@@ -0,0 +1,581 @@
+"""
+Cloud Data Sources - BigQuery Integration
+Tools for loading and writing data to/from Google BigQuery.
+Compatible with existing DataScienceCopilot tool registry.
+"""
+
+import polars as pl
+import pandas as pd
+from typing import Dict, Any, Optional, Literal
+from pathlib import Path
+import sys
+import os
+
+# Add parent directory to path
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from ds_agent.utils.validation import validate_dataframe
+
+try:
+    from google.cloud import bigquery
+    from google.oauth2 import service_account
+    BIGQUERY_AVAILABLE = True
+except ImportError:
+    BIGQUERY_AVAILABLE = False
+    bigquery = None
+    service_account = None
+
+
+def _get_bigquery_client(project_id: str) -> 'bigquery.Client':
+    """
+    Initialize BigQuery client with credentials from environment.
+    
+    Credential sources (in order of priority):
+    1. GOOGLE_APPLICATION_CREDENTIALS env var (service account JSON path)
+    2. Default application credentials (gcloud auth application-default login)
+    
+    Args:
+        project_id: Google Cloud project ID
+        
+    Returns:
+        BigQuery client instance
+        
+    Raises:
+        ImportError: If google-cloud-bigquery not installed
+        EnvironmentError: If credentials not found
+    """
+    if not BIGQUERY_AVAILABLE:
+        raise ImportError(
+            "google-cloud-bigquery is not installed. "
+            "Install it with: pip install google-cloud-bigquery"
+        )
+    
+    # Check for service account credentials
+    creds_path = os.getenv("GOOGLE_APPLICATION_CREDENTIALS")
+    
+    if creds_path and Path(creds_path).exists():
+        # Use service account JSON
+        credentials = service_account.Credentials.from_service_account_file(creds_path)
+        client = bigquery.Client(project=project_id, credentials=credentials)
+    else:
+        # Use default application credentials
+        try:
+            client = bigquery.Client(project=project_id)
+        except Exception as e:
+            raise EnvironmentError(
+                "BigQuery credentials not found. Either:\n"
+                "1. Set GOOGLE_APPLICATION_CREDENTIALS to service account JSON path\n"
+                "2. Run: gcloud auth application-default login\n"
+                f"Error: {str(e)}"
+            )
+    
+    return client
+
+
+def load_bigquery_table(
+    project_id: str,
+    dataset: str,
+    table: str,
+    limit: Optional[int] = None,
+    columns: Optional[list] = None,
+    where_clause: Optional[str] = None
+) -> Dict[str, Any]:
+    """
+    Load data from BigQuery table into a Polars DataFrame.
+    
+    This tool allows the agent to load data from BigQuery for analysis.
+    Supports sampling via LIMIT and column selection for memory efficiency.
+    
+    Args:
+        project_id: Google Cloud project ID
+        dataset: BigQuery dataset name
+        table: BigQuery table name
+        limit: Optional row limit for sampling (e.g., 10000 for large tables)
+        columns: Optional list of column names to load (default: all columns)
+        where_clause: Optional SQL WHERE clause for filtering (without WHERE keyword)
+            Example: "created_at > '2024-01-01'"
+    
+    Returns:
+        Dictionary with:
+        - success: bool
+        - data_path: str (saved CSV path for downstream tools)
+        - df_info: dict (shape, columns, memory_usage)
+        - message: str
+        - query_stats: dict (bytes processed, rows returned)
+    
+    Examples:
+        >>> # Load full table
+        >>> load_bigquery_table("my-project", "analytics", "users")
+        
+        >>> # Sample 10K rows for exploration
+        >>> load_bigquery_table("my-project", "analytics", "events", limit=10000)
+        
+        >>> # Load specific columns with filter
+        >>> load_bigquery_table(
+        ...     "my-project", "sales", "transactions",
+        ...     columns=["customer_id", "amount", "date"],
+        ...     where_clause="date >= '2024-01-01'",
+        ...     limit=50000
+        ... )
+    """
+    try:
+        # Initialize client
+        client = _get_bigquery_client(project_id)
+        
+        # Build query
+        table_ref = f"{project_id}.{dataset}.{table}"
+        
+        if columns:
+            columns_str = ", ".join(columns)
+        else:
+            columns_str = "*"
+        
+        query = f"SELECT {columns_str} FROM `{table_ref}`"
+        
+        if where_clause:
+            query += f" WHERE {where_clause}"
+        
+        if limit:
+            query += f" LIMIT {limit}"
+        
+        # Execute query
+        query_job = client.query(query)
+        
+        # Load results into pandas (BigQuery SDK returns pandas)
+        df_pandas = query_job.to_dataframe()
+        
+        # Convert to Polars for consistency with existing tools
+        df = pl.from_pandas(df_pandas)
+        
+        # Validate
+        validate_dataframe(df)
+        
+        # Save to outputs/data/ for downstream tool compatibility
+        output_dir = Path("./outputs/data")
+        output_dir.mkdir(parents=True, exist_ok=True)
+        
+        output_path = output_dir / f"bigquery_{dataset}_{table}.csv"
+        df.write_csv(output_path)
+        
+        # Get query statistics
+        bytes_processed = query_job.total_bytes_processed or 0
+        bytes_billed = query_job.total_bytes_billed or 0
+        
+        return {
+            "success": True,
+            "data_path": str(output_path),
+            "df_info": {
+                "rows": df.shape[0],
+                "columns": df.shape[1],
+                "column_names": df.columns,
+                "memory_mb": round(df.estimated_size("mb"), 2)
+            },
+            "query_stats": {
+                "bytes_processed": bytes_processed,
+                "bytes_processed_mb": round(bytes_processed / 1024 / 1024, 2),
+                "bytes_billed": bytes_billed,
+                "bytes_billed_mb": round(bytes_billed / 1024 / 1024, 2),
+                "rows_returned": len(df)
+            },
+            "message": f"✅ Loaded {len(df):,} rows from {table_ref}. Saved to {output_path}",
+            "table_reference": table_ref,
+            "query": query
+        }
+    
+    except ImportError as e:
+        return {
+            "success": False,
+            "error": str(e),
+            "error_type": "ImportError",
+            "message": "BigQuery library not installed. Run: pip install google-cloud-bigquery"
+        }
+    
+    except Exception as e:
+        return {
+            "success": False,
+            "error": str(e),
+            "error_type": type(e).__name__,
+            "message": f"Failed to load BigQuery table: {str(e)}"
+        }
+
+
+def write_bigquery_table(
+    file_path: str,
+    project_id: str,
+    dataset: str,
+    table: str,
+    mode: Literal["append", "overwrite", "fail"] = "append"
+) -> Dict[str, Any]:
+    """
+    Write DataFrame to BigQuery table from CSV/Parquet file.
+    
+    This tool allows the agent to save predictions, metrics, or processed data
+    back to BigQuery for downstream consumption.
+    
+    Args:
+        file_path: Path to CSV or Parquet file containing data to write
+        project_id: Google Cloud project ID
+        dataset: BigQuery dataset name
+        table: BigQuery table name
+        mode: Write mode
+            - "append": Add rows to existing table
+            - "overwrite": Replace table contents
+            - "fail": Raise error if table exists
+    
+    Returns:
+        Dictionary with:
+        - success: bool
+        - table_reference: str
+        - rows_written: int
+        - message: str
+    
+    Examples:
+        >>> # Write predictions to BigQuery
+        >>> write_bigquery_table(
+        ...     "./outputs/data/predictions.csv",
+        ...     "my-project",
+        ...     "ml_results",
+        ...     "churn_predictions",
+        ...     mode="append"
+        ... )
+        
+        >>> # Overwrite existing metrics table
+        >>> write_bigquery_table(
+        ...     "./outputs/data/metrics.csv",
+        ...     "my-project",
+        ...     "ml_results",
+        ...     "model_metrics",
+        ...     mode="overwrite"
+        ... )
+    """
+    try:
+        # Initialize client
+        client = _get_bigquery_client(project_id)
+        
+        # Load data from file
+        file_path = Path(file_path)
+        if not file_path.exists():
+            return {
+                "success": False,
+                "error": f"File not found: {file_path}",
+                "error_type": "FileNotFoundError"
+            }
+        
+        # Load based on extension
+        if file_path.suffix.lower() == ".csv":
+            df = pl.read_csv(file_path)
+        elif file_path.suffix.lower() == ".parquet":
+            df = pl.read_parquet(file_path)
+        else:
+            return {
+                "success": False,
+                "error": f"Unsupported file format: {file_path.suffix}",
+                "error_type": "ValueError"
+            }
+        
+        # Convert to pandas (BigQuery SDK requires pandas)
+        df_pandas = df.to_pandas()
+        
+        # Build table reference
+        table_ref = f"{project_id}.{dataset}.{table}"
+        
+        # Configure write disposition
+        if mode == "append":
+            write_disposition = bigquery.WriteDisposition.WRITE_APPEND
+        elif mode == "overwrite":
+            write_disposition = bigquery.WriteDisposition.WRITE_TRUNCATE
+        elif mode == "fail":
+            write_disposition = bigquery.WriteDisposition.WRITE_EMPTY
+        else:
+            return {
+                "success": False,
+                "error": f"Invalid mode: {mode}. Use 'append', 'overwrite', or 'fail'",
+                "error_type": "ValueError"
+            }
+        
+        # Configure job
+        job_config = bigquery.LoadJobConfig(
+            write_disposition=write_disposition,
+            autodetect=True  # Auto-detect schema from DataFrame
+        )
+        
+        # Execute write job
+        job = client.load_table_from_dataframe(
+            df_pandas,
+            table_ref,
+            job_config=job_config
+        )
+        
+        # Wait for completion
+        job.result()
+        
+        return {
+            "success": True,
+            "table_reference": table_ref,
+            "rows_written": len(df_pandas),
+            "mode": mode,
+            "message": f"✅ Wrote {len(df_pandas):,} rows to {table_ref} (mode: {mode})",
+            "table_info": {
+                "project": project_id,
+                "dataset": dataset,
+                "table": table,
+                "columns": df.columns,
+                "rows": len(df)
+            }
+        }
+    
+    except ImportError as e:
+        return {
+            "success": False,
+            "error": str(e),
+            "error_type": "ImportError",
+            "message": "BigQuery library not installed. Run: pip install google-cloud-bigquery"
+        }
+    
+    except Exception as e:
+        return {
+            "success": False,
+            "error": str(e),
+            "error_type": type(e).__name__,
+            "message": f"Failed to write to BigQuery: {str(e)}"
+        }
+
+
+def profile_bigquery_table(
+    project_id: str,
+    dataset: str,
+    table: str
+) -> Dict[str, Any]:
+    """
+    Profile a BigQuery table without loading all data.
+    
+    Returns metadata including row count, column types, null counts,
+    and table size. Useful for initial exploration before full load.
+    
+    Args:
+        project_id: Google Cloud project ID
+        dataset: BigQuery dataset name
+        table: BigQuery table name
+    
+    Returns:
+        Dictionary with:
+        - success: bool
+        - table_reference: str
+        - row_count: int
+        - columns: list of dicts with column info
+        - table_size_mb: float
+        - created: str (timestamp)
+        - modified: str (timestamp)
+        - message: str
+    
+    Examples:
+        >>> # Quick profile before loading
+        >>> profile_bigquery_table("my-project", "analytics", "events")
+        {
+            "success": True,
+            "row_count": 1000000,
+            "columns": [
+                {"name": "user_id", "type": "STRING", "mode": "NULLABLE"},
+                {"name": "event_time", "type": "TIMESTAMP", "mode": "REQUIRED"},
+                ...
+            ],
+            "table_size_mb": 125.5
+        }
+    """
+    try:
+        # Initialize client
+        client = _get_bigquery_client(project_id)
+        
+        # Get table metadata
+        table_ref = f"{project_id}.{dataset}.{table}"
+        table_obj = client.get_table(table_ref)
+        
+        # Extract schema information
+        columns_info = []
+        for field in table_obj.schema:
+            columns_info.append({
+                "name": field.name,
+                "type": field.field_type,
+                "mode": field.mode,  # NULLABLE, REQUIRED, REPEATED
+                "description": field.description or ""
+            })
+        
+        # Get null counts via query (sample for efficiency)
+        null_counts = {}
+        try:
+            # Use TABLESAMPLE for large tables (1% sample)
+            sample_query = f"""
+            SELECT 
+                {', '.join([f'COUNTIF({col["name"]} IS NULL) AS {col["name"]}_nulls' for col in columns_info])}
+            FROM `{table_ref}`
+            TABLESAMPLE SYSTEM (1 PERCENT)
+            """
+            
+            query_job = client.query(sample_query)
+            result = query_job.result()
+            row = next(iter(result))
+            
+            for col in columns_info:
+                null_count = row.get(f'{col["name"]}_nulls', 0)
+                null_counts[col["name"]] = null_count
+        except Exception as e:
+            # If sampling fails, skip null counts
+            null_counts = {col["name"]: "N/A" for col in columns_info}
+        
+        # Table size information
+        table_size_bytes = table_obj.num_bytes or 0
+        table_size_mb = round(table_size_bytes / 1024 / 1024, 2)
+        
+        return {
+            "success": True,
+            "table_reference": table_ref,
+            "profile": {
+                "row_count": table_obj.num_rows,
+                "column_count": len(columns_info),
+                "table_size_mb": table_size_mb,
+                "table_size_gb": round(table_size_mb / 1024, 2)
+            },
+            "columns": columns_info,
+            "null_counts_sample": null_counts,
+            "metadata": {
+                "created": table_obj.created.isoformat() if table_obj.created else None,
+                "modified": table_obj.modified.isoformat() if table_obj.modified else None,
+                "location": table_obj.location,
+                "expiration": table_obj.expires.isoformat() if table_obj.expires else None
+            },
+            "message": f"✅ Profiled {table_ref}: {table_obj.num_rows:,} rows, {len(columns_info)} columns, {table_size_mb} MB",
+            "recommendation": (
+                f"Table has {table_obj.num_rows:,} rows. "
+                f"Consider using limit={min(10000, table_obj.num_rows)} for initial exploration."
+                if table_obj.num_rows > 10000 else
+                f"Table is small ({table_obj.num_rows:,} rows), safe to load fully."
+            )
+        }
+    
+    except ImportError as e:
+        return {
+            "success": False,
+            "error": str(e),
+            "error_type": "ImportError",
+            "message": "BigQuery library not installed. Run: pip install google-cloud-bigquery"
+        }
+    
+    except Exception as e:
+        return {
+            "success": False,
+            "error": str(e),
+            "error_type": type(e).__name__,
+            "message": f"Failed to profile BigQuery table: {str(e)}"
+        }
+
+
+def query_bigquery(
+    project_id: str,
+    query: str,
+    output_path: Optional[str] = None,
+    limit: Optional[int] = None
+) -> Dict[str, Any]:
+    """
+    Execute a custom BigQuery SQL query and return results as DataFrame.
+    
+    This tool allows the agent to run custom SQL queries for complex
+    data transformations before analysis.
+    
+    Args:
+        project_id: Google Cloud project ID
+        query: SQL query to execute
+        output_path: Optional path to save results (default: auto-generated)
+        limit: Optional row limit to append to query
+    
+    Returns:
+        Dictionary with:
+        - success: bool
+        - data_path: str
+        - df_info: dict
+        - query_stats: dict
+        - message: str
+    
+    Examples:
+        >>> # Custom aggregation query
+        >>> query_bigquery(
+        ...     "my-project",
+        ...     '''
+        ...     SELECT 
+        ...         customer_id,
+        ...         SUM(amount) as total_spent,
+        ...         COUNT(*) as num_orders
+        ...     FROM `my-project.sales.orders`
+        ...     WHERE date >= '2024-01-01'
+        ...     GROUP BY customer_id
+        ...     '''
+        ... )
+    """
+    try:
+        # Initialize client
+        client = _get_bigquery_client(project_id)
+        
+        # Add limit if specified
+        if limit:
+            query = f"{query.rstrip(';')} LIMIT {limit}"
+        
+        # Execute query
+        query_job = client.query(query)
+        df_pandas = query_job.to_dataframe()
+        
+        # Convert to Polars
+        df = pl.from_pandas(df_pandas)
+        
+        # Determine output path
+        if output_path is None:
+            output_dir = Path("./outputs/data")
+            output_dir.mkdir(parents=True, exist_ok=True)
+            output_path = str(output_dir / "bigquery_query_result.csv")
+        
+        # Save results
+        df.write_csv(output_path)
+        
+        # Get query statistics
+        bytes_processed = query_job.total_bytes_processed or 0
+        
+        return {
+            "success": True,
+            "data_path": output_path,
+            "df_info": {
+                "rows": df.shape[0],
+                "columns": df.shape[1],
+                "column_names": df.columns,
+                "memory_mb": round(df.estimated_size("mb"), 2)
+            },
+            "query_stats": {
+                "bytes_processed": bytes_processed,
+                "bytes_processed_mb": round(bytes_processed / 1024 / 1024, 2),
+                "rows_returned": len(df)
+            },
+            "message": f"✅ Query returned {len(df):,} rows. Saved to {output_path}",
+            "query": query
+        }
+    
+    except ImportError as e:
+        return {
+            "success": False,
+            "error": str(e),
+            "error_type": "ImportError",
+            "message": "BigQuery library not installed. Run: pip install google-cloud-bigquery"
+        }
+    
+    except Exception as e:
+        return {
+            "success": False,
+            "error": str(e),
+            "error_type": type(e).__name__,
+            "message": f"Failed to execute BigQuery query: {str(e)}"
+        }
+
+
+# Export functions for tool registry
+__all__ = [
+    'load_bigquery_table',
+    'write_bigquery_table',
+    'profile_bigquery_table',
+    'query_bigquery'
+]