dbagent-cli 0.1.0__py3-none-any.whl

dbagent/config.py

@@ -0,0 +1,89 @@
+"""
+Configuration and database profile manager.
+Stores profiles and credentials in user home directory (~/.dbagent/profiles.json).
+"""
+
+import json
+import os
+from pathlib import Path
+from typing import Dict, Any, Optional
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+class ConfigManager:
+    """Manages local profiles and AI API credentials."""
+
+    def __init__(self, config_dir: Optional[Path] = None):
+        self.config_dir = config_dir or Path.home() / ".dbagent"
+        self.config_file = self.config_dir / "config.json"
+        self.profiles_file = self.config_dir / "profiles.json"
+        self._ensure_dirs()
+
+    def _ensure_dirs(self) -> None:
+        self.config_dir.mkdir(parents=True, exist_ok=True)
+        if not self.config_file.exists():
+            self._save_json(self.config_file, {
+                "default_provider": "ollama",
+                "default_model": None,
+                "ollama_base_url": "http://localhost:11434",
+                "gemini_api_key": None,
+                "groq_api_key": None,
+                "openrouter_api_key": None,
+            })
+        if not self.profiles_file.exists():
+            self._save_json(self.profiles_file, {})
+
+    def _load_json(self, path: Path) -> Dict[str, Any]:
+        try:
+            with open(path, "r", encoding="utf-8") as f:
+                return json.load(f)
+        except Exception:
+            return {}
+
+    def _save_json(self, path: Path, data: Dict[str, Any]) -> None:
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump(data, f, indent=2)
+
+    # --- LLM Settings ---
+
+    def get_setting(self, key: str, default: Any = None) -> Any:
+        # Check environment variable first
+        env_key = key.upper()
+        if os.getenv(env_key):
+            return os.getenv(env_key)
+
+        config = self._load_json(self.config_file)
+        return config.get(key, default)
+
+    def set_setting(self, key: str, value: Any) -> None:
+        config = self._load_json(self.config_file)
+        config[key] = value
+        self._save_json(self.config_file, config)
+
+    # --- Database Profiles ---
+
+    def save_profile(self, name: str, connection_url: str, description: Optional[str] = None) -> None:
+        profiles = self._load_json(self.profiles_file)
+        profiles[name] = {
+            "name": name,
+            "url": connection_url,
+            "description": description or f"Profile for {name}",
+        }
+        self._save_json(self.profiles_file, profiles)
+
+    def get_profile(self, name: str) -> Optional[Dict[str, Any]]:
+        profiles = self._load_json(self.profiles_file)
+        return profiles.get(name)
+
+    def list_profiles(self) -> Dict[str, Dict[str, Any]]:
+        return self._load_json(self.profiles_file)
+
+    def delete_profile(self, name: str) -> bool:
+        profiles = self._load_json(self.profiles_file)
+        if name in profiles:
+            del profiles[name]
+            self._save_json(self.profiles_file, profiles)
+            return True
+        return False

dbagent/connectors/base.py

@@ -0,0 +1,75 @@
+"""
+Base connector interface for database introspection and query execution.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Dict, Any, List, Optional, Tuple, Set
+from dbagent.schema.models import DatabaseSchema
+
+
+class BaseConnector(ABC):
+    """Abstract interface for database connectors."""
+
+    @abstractmethod
+    def test_connection(self) -> Tuple[bool, str]:
+        """Test if the database connection is working. Returns (success, message)."""
+        pass
+
+    @abstractmethod
+    def get_table_names(self) -> List[str]:
+        """Get fast list of all table names without full column introspection."""
+        pass
+
+    @abstractmethod
+    def inspect_schema(
+        self,
+        table_names: Optional[List[str]] = None,
+        include_samples: bool = True,
+        max_samples: int = 2,
+        include_views: bool = True,
+        include_row_counts: bool = False,
+    ) -> DatabaseSchema:
+        """Introspect the database schema for all or specified tables."""
+        pass
+
+    @abstractmethod
+    def inspect_targeted(
+        self,
+        user_prompt: str,
+        max_tables: int = 8,
+        include_samples: bool = True,
+    ) -> DatabaseSchema:
+        """Introspect only the tables relevant to the user query on demand."""
+        pass
+
+    def resolve_tables(self, user_prompt: str) -> Tuple[List[str], List[str]]:
+        """
+        Resolve which tables the user is referring to.
+        Returns (exact_matches, fuzzy_matches).
+        Default implementation uses get_table_names with simple substring matching.
+        """
+        all_tables = self.get_table_names()
+        prompt_lower = user_prompt.lower()
+        exact = [t for t in all_tables if t.lower() in prompt_lower]
+        return exact, []
+
+    def refresh_cache(self) -> None:
+        """Invalidate any cached metadata. Override in subclasses."""
+        pass
+
+    @abstractmethod
+    def execute_query(
+        self,
+        query: str,
+        limit: int = 100,
+    ) -> Tuple[List[str], List[Dict[str, Any]], Optional[str]]:
+        """
+        Execute a read query.
+        Returns (columns, rows, error_message).
+        """
+        pass
+
+    @abstractmethod
+    def close(self) -> None:
+        """Close any open connections."""
+        pass

dbagent/connectors/factory.py

@@ -0,0 +1,25 @@
+"""
+Connector factory for instantiating the right database connector.
+"""
+
+from dbagent.connectors.base import BaseConnector
+from dbagent.connectors.relational import RelationalConnector
+
+
+def create_connector(connection_url: str, **kwargs) -> BaseConnector:
+    """Create a database connector based on the connection string URL."""
+    url = connection_url.strip()
+
+    # Normalize SQLite file paths
+    if url.endswith(".db") or url.endswith(".sqlite") or url.endswith(".sqlite3"):
+        if not url.startswith("sqlite:///"):
+            url = f"sqlite:///{url}"
+
+    # MongoDB
+    if url.startswith("mongodb://") or url.startswith("mongodb+srv://"):
+        from dbagent.connectors.mongo import MongoConnector
+        db_name = kwargs.get("database") or kwargs.get("db_name")
+        return MongoConnector(connection_url=url, db_name=db_name)
+
+    # All relational databases
+    return RelationalConnector(connection_url=url)

dbagent/connectors/mongo.py

@@ -0,0 +1,237 @@
+"""
+MongoDB Connector for document database schema introspection and query execution.
+"""
+
+from typing import Dict, Any, List, Optional, Tuple, Set
+from dbagent.connectors.base import BaseConnector
+from dbagent.schema.models import (
+    DatabaseSchema,
+    TableModel,
+    ColumnModel,
+    IndexModel,
+)
+
+
+def _infer_bson_type(val: Any) -> str:
+    """Infer BSON type as a human-readable string."""
+    if val is None:
+        return "Null"
+    if isinstance(val, bool):
+        return "Boolean"
+    if isinstance(val, int):
+        return "Integer"
+    if isinstance(val, float):
+        return "Double"
+    if isinstance(val, str):
+        return "String"
+    if isinstance(val, list):
+        if val:
+            inner_type = _infer_bson_type(val[0])
+            return f"Array<{inner_type}>"
+        return "Array"
+    if isinstance(val, dict):
+        return "Object"
+    return type(val).__name__
+
+
+class MongoConnector(BaseConnector):
+    """Connector for MongoDB instances."""
+
+    def __init__(self, connection_url: str, db_name: Optional[str] = None):
+        self.connection_url = connection_url
+        self.explicit_db_name = db_name
+        self._client = None
+        self._cached_collection_names: Optional[List[str]] = None
+
+    def _get_client(self):
+        if self._client is None:
+            import pymongo
+            self._client = pymongo.MongoClient(self.connection_url, serverSelectionTimeoutMS=5000)
+        return self._client
+
+    def _get_database(self):
+        client = self._get_client()
+        if self.explicit_db_name:
+            return client[self.explicit_db_name]
+        default_db = client.get_default_database(default=None)
+        if default_db is not None:
+            return default_db
+        # Fall back to first non-admin/non-local db
+        db_names = [d for d in client.list_database_names() if d not in ["admin", "config", "local"]]
+        if db_names:
+            return client[db_names[0]]
+        return client["test"]
+
+    def test_connection(self) -> Tuple[bool, str]:
+        try:
+            client = self._get_client()
+            client.admin.command("ping")
+            db = self._get_database()
+            return True, f"Successfully connected to MongoDB (database: '{db.name}')."
+        except Exception as e:
+            return False, f"MongoDB connection failed: {str(e)}"
+
+    def get_table_names(self) -> List[str]:
+        if self._cached_collection_names is not None:
+            return self._cached_collection_names
+        db = self._get_database()
+        self._cached_collection_names = db.list_collection_names()
+        return self._cached_collection_names
+
+    def refresh_cache(self) -> None:
+        """Invalidate cached collection names."""
+        self._cached_collection_names = None
+
+    def resolve_tables(self, user_prompt: str) -> Tuple[List[str], List[str]]:
+        """Resolve which collections the user is referring to with fuzzy matching."""
+        from dbagent.connectors.relational import _fuzzy_match_tables
+        all_cols = self.get_table_names()
+        return _fuzzy_match_tables(all_cols, user_prompt)
+
+    def inspect_schema(
+        self,
+        table_names: Optional[List[str]] = None,
+        include_samples: bool = True,
+        max_samples: int = 2,
+        include_views: bool = True,
+        include_row_counts: bool = False,
+    ) -> DatabaseSchema:
+        """Introspect MongoDB collections by sampling documents and inspecting indexes."""
+        db = self._get_database()
+        collection_names = db.list_collection_names()
+
+        target_collections = collection_names
+        if table_names is not None:
+            lower_targets = {t.lower() for t in table_names}
+            target_collections = [c for c in collection_names if c.lower() in lower_targets]
+
+        tables: List[TableModel] = []
+
+        for c_name in target_collections:
+            col = db[c_name]
+            doc_count = col.estimated_document_count() if include_row_counts else None
+
+            # Sample documents to infer schema
+            sample_docs = list(col.find().limit(max(max_samples, 10)))
+            fields_map: Dict[str, set] = {}
+            for doc in sample_docs:
+                for k, v in doc.items():
+                    if k not in fields_map:
+                        fields_map[k] = set()
+                    fields_map[k].add(_infer_bson_type(v))
+
+            columns: List[ColumnModel] = []
+            for field_name, types in fields_map.items():
+                type_str = " | ".join(sorted(types)) if types else "Unknown"
+                is_pk = field_name == "_id"
+                columns.append(
+                    ColumnModel(
+                        name=field_name,
+                        data_type=type_str,
+                        is_nullable=not is_pk,
+                        is_primary_key=is_pk,
+                    )
+                )
+
+            # Indexes
+            indexes: List[IndexModel] = []
+            try:
+                raw_indexes = col.index_information()
+                for idx_name, idx_info in raw_indexes.items():
+                    key_cols = [k[0] for k in idx_info.get("key", [])]
+                    indexes.append(
+                        IndexModel(
+                            name=idx_name,
+                            columns=key_cols,
+                            is_unique=idx_info.get("unique", False),
+                        )
+                    )
+            except Exception:
+                pass
+
+            # Serialize sample rows
+            clean_samples = []
+            if include_samples:
+                for doc in sample_docs[:max_samples]:
+                    clean_doc = {}
+                    for k, v in doc.items():
+                        clean_doc[k] = str(v) if not isinstance(v, (int, float, bool, str)) else v
+                    clean_samples.append(clean_doc)
+
+            tables.append(
+                TableModel(
+                    name=c_name,
+                    is_view=False,
+                    columns=columns,
+                    primary_key=["_id"],
+                    foreign_keys=[],
+                    indexes=indexes,
+                    row_count=doc_count,
+                    sample_rows=clean_samples,
+                )
+            )
+
+        return DatabaseSchema(
+            dialect_name="mongodb",
+            database_name=db.name,
+            tables=tables,
+        )
+
+    def inspect_targeted(
+        self,
+        user_prompt: str,
+        max_tables: int = 8,
+        include_samples: bool = True,
+    ) -> DatabaseSchema:
+        """On-demand targeted MongoDB collection inspection."""
+        import re
+        all_cols = self.get_table_names()
+        prompt_lower = user_prompt.lower()
+
+        matched = []
+        for c in all_cols:
+            c_lower = c.lower()
+            if re.search(r"\b" + re.escape(c_lower) + r"\b", prompt_lower) or c_lower in prompt_lower:
+                matched.append(c)
+
+        if matched:
+            return self.inspect_schema(
+                table_names=matched[:max_tables],
+                include_samples=include_samples,
+                max_samples=2,
+                include_row_counts=False,
+            )
+
+        return self.inspect_schema(
+            table_names=all_cols[:max_tables],
+            include_samples=include_samples,
+            max_samples=2,
+            include_row_counts=False,
+        )
+
+    def execute_query(
+        self,
+        query: str,
+        limit: int = 100,
+    ) -> Tuple[List[str], List[Dict[str, Any]], Optional[str]]:
+        """For MongoDB, queries can be formatted as JSON filter or JS collection queries."""
+        try:
+            import json
+            # If query is JSON formatted e.g. {"collection": "users", "filter": {}}
+            parsed = json.loads(query)
+            col_name = parsed.get("collection")
+            filt = parsed.get("filter", {})
+            db = self._get_database()
+            cursor = db[col_name].find(filt).limit(limit)
+            docs = list(cursor)
+            if not docs:
+                return [], [], None
+            columns = list(docs[0].keys())
+            rows = [{k: str(v) for k, v in d.items()} for d in docs]
+            return columns, rows, None
+        except Exception as e:
+            return [], [], f"MongoDB execution error: {str(e)}"
+
+    def close(self) -> None:
+        if self._client:
+            self._client.close()