mcp-server-motherduck

import os
import duckdb
from typing import Literal, Optional
from tabulate import tabulate
import logging
import threading
from .configs import SERVER_VERSION

logger = logging.getLogger("mcp_server_motherduck")


class DatabaseClient:
    def __init__(
        self,
        db_path: str | None = None,
        motherduck_token: str | None = None,
        home_dir: str | None = None,
        saas_mode: bool = False,
        read_only: bool = False,
        max_rows: int = 1024,
        max_chars: int = 50000,
        query_timeout: int = -1,
    ):
        self._read_only = read_only
        self._max_rows = max_rows
        self._max_chars = max_chars
        self._query_timeout = query_timeout
        self.db_path, self.db_type = self._resolve_db_path_type(
            db_path, motherduck_token, saas_mode
        )
        logger.info(f"Database client initialized in `{self.db_type}` mode")

        # Set the home directory for DuckDB
        if home_dir:
            os.environ["HOME"] = home_dir

        self.conn = self._initialize_connection()

    def _initialize_connection(self) -> Optional[duckdb.DuckDBPyConnection]:
        """Initialize connection to the MotherDuck or DuckDB database"""

        logger.info(f"🔌 Connecting to {self.db_type} database")

        # S3 databases don't support read-only mode
        if self.db_type == "s3" and self._read_only:
            raise ValueError("Read-only mode is not supported for S3 databases")

        if self.db_type == "duckdb" and self._read_only:
            # check that we can connect, issue a `select 1` and then close + return None
            try:
                conn = duckdb.connect(
                    self.db_path,
                    config={
                        "custom_user_agent": f"mcp-server-motherduck/{SERVER_VERSION}"
                    },
                    read_only=self._read_only,
                )
                conn.execute("SELECT 1")
                conn.close()
                return None
            except Exception as e:
                logger.error(f"❌ Read-only check failed: {e}")
                raise

        # Check if this is an S3 path
        if self.db_type == "s3":
            # For S3, we need to create an in-memory connection and attach the S3 database
            conn = duckdb.connect(':memory:')
            
            # Install and load the httpfs extension for S3 support
            import io
            from contextlib import redirect_stdout, redirect_stderr
            
            null_file = io.StringIO()
            with redirect_stdout(null_file), redirect_stderr(null_file):
                try:
                    conn.execute("INSTALL httpfs;")
                except:
                    pass  # Extension might already be installed
                conn.execute("LOAD httpfs;")
            
            # Configure S3 credentials from environment variables using CREATE SECRET
            aws_access_key = os.environ.get('AWS_ACCESS_KEY_ID')
            aws_secret_key = os.environ.get('AWS_SECRET_ACCESS_KEY')
            aws_session_token = os.environ.get('AWS_SESSION_TOKEN')
            aws_region = os.environ.get('AWS_DEFAULT_REGION', 'us-east-1')
            
            
            if aws_access_key and aws_secret_key and not aws_session_token:
                # Use CREATE SECRET for better credential management
                conn.execute(f"""
                    CREATE SECRET IF NOT EXISTS s3_secret (
                        TYPE S3,
                        KEY_ID '{aws_access_key}',
                        SECRET '{aws_secret_key}',
                        REGION '{aws_region}'
                    );
                """)
            elif aws_session_token:
                # Use credential_chain provider to automatically fetch credentials
                # This supports IAM roles, SSO, instance profiles, etc.
                conn.execute(f"""
                    CREATE SECRET IF NOT EXISTS s3_secret (
                        TYPE S3,
                        PROVIDER credential_chain,
                        REGION '{aws_region}'
                    );
                """)
            
            # Attach the S3 database
            try:
                # For S3, we always attach as READ_ONLY since S3 storage is typically read-only
                # Even when not in read_only mode, we attach as READ_ONLY for S3
                conn.execute(f"ATTACH '{self.db_path}' AS s3db (READ_ONLY);")
                # Use the attached database
                conn.execute("USE s3db;")
                logger.info(f"✅ Successfully connected to {self.db_type} database (attached as read-only)")
            except Exception as e:
                logger.error(f"Failed to attach S3 database: {e}")
                # If the database doesn't exist and we're not in read-only mode, try to create it
                if "database does not exist" in str(e) and not self._read_only:
                    logger.info("S3 database doesn't exist, attempting to create it...")
                    try:
                        # Create a new database at the S3 location
                        conn.execute(f"ATTACH '{self.db_path}' AS s3db;")
                        conn.execute("USE s3db;")
                        logger.info(f"✅ Created new S3 database at {self.db_path}")
                    except Exception as create_error:
                        logger.error(f"Failed to create S3 database: {create_error}")
                        raise
                else:
                    raise
                
            return conn

        conn = duckdb.connect(
            self.db_path,
            config={"custom_user_agent": f"mcp-server-motherduck/{SERVER_VERSION}"},
            read_only=self._read_only,
        )

        logger.info(f"✅ Successfully connected to {self.db_type} database")

        return conn

    def _resolve_db_path_type(
        self, db_path: str, motherduck_token: str | None = None, saas_mode: bool = False
    ) -> tuple[str, Literal["duckdb", "motherduck", "s3"]]:
        """Resolve and validate the database path"""
        # Handle S3 paths
        if db_path.startswith("s3://"):
            return db_path, "s3"
        
        # Handle MotherDuck paths
        if db_path.startswith("md:"):
            if motherduck_token:
                logger.info("Using MotherDuck token to connect to database `md:`")
                if saas_mode:
                    logger.info("Connecting to MotherDuck in SaaS mode")
                    return (
                        f"{db_path}?motherduck_token={motherduck_token}&saas_mode=true",
                        "motherduck",
                    )
                else:
                    return (
                        f"{db_path}?motherduck_token={motherduck_token}",
                        "motherduck",
                    )
            elif os.getenv("motherduck_token"):
                logger.info(
                    "Using MotherDuck token from env to connect to database `md:`"
                )
                return (
                    f"{db_path}?motherduck_token={os.getenv('motherduck_token')}",
                    "motherduck",
                )
            else:
                raise ValueError(
                    "Please set the `motherduck_token` as an environment variable or pass it as an argument with `--motherduck-token` when using `md:` as db_path."
                )

        if db_path == ":memory:":
            return db_path, "duckdb"

        return db_path, "duckdb"

    def _execute(self, query: str) -> str:
        # Get connection to use
        if self.conn is None:
            conn = duckdb.connect(
                self.db_path,
                config={"custom_user_agent": f"mcp-server-motherduck/{SERVER_VERSION}"},
                read_only=self._read_only,
            )
        else:
            conn = self.conn
        
        # Execute with or without timeout
        if self._query_timeout > 0:
            rows, has_more_rows, headers = self._execute_with_timeout(conn, query)
        else:
            rows, has_more_rows, headers = self._execute_direct(conn, query)
        
        # Close connection if it was temporary
        if self.conn is None:
            conn.close()
        
        returned_rows = len(rows)
        
        # Format results as table
        out = tabulate(rows, headers=headers, tablefmt="pretty")
        
        # Apply character limit if output is too long
        char_truncated = len(out) > self._max_chars
        if char_truncated:
            out = out[:self._max_chars]
        
        # Add informative feedback message
        if has_more_rows:
            out += f"\n\n⚠️  Showing first {returned_rows} rows."
        elif char_truncated:
            out += f"\n\n⚠️  Output truncated at {self._max_chars:,} characters."

        return out
    
    def _execute_direct(self, conn, query: str) -> tuple:
        """Execute query without timeout - original code path"""
        q = conn.execute(query)
        rows = q.fetchmany(self._max_rows)
        has_more_rows = q.fetchone() is not None
        headers = [d[0] + "\n" + str(d[1]) for d in q.description]
        return rows, has_more_rows, headers
    
    def _execute_with_timeout(self, conn, query: str) -> tuple:
        """Execute query with timeout using threading.Timer and conn.interrupt()"""
        timer = threading.Timer(self._query_timeout, conn.interrupt)
        timer.start()
        
        try:
            q = conn.execute(query)
            rows = q.fetchmany(self._max_rows)
            has_more_rows = q.fetchone() is not None
            headers = [d[0] + "\n" + str(d[1]) for d in q.description]
            return rows, has_more_rows, headers
        except duckdb.InterruptException:
            raise ValueError(
                f"Query execution timed out after {self._query_timeout} seconds. "
                f"Increase timeout with --query-timeout argument when starting the mcp server."
            )
        finally:
            timer.cancel()

    def query(self, query: str) -> str:
        try:
            return self._execute(query)

        except Exception as e:
            raise ValueError(f"❌ Error executing query: {e}")