db-connect-mcp 0.1.0__py3-none-any.whl

db_connect_mcp/models/statistics.py

@@ -0,0 +1,176 @@
+"""Column statistics and distribution models."""
+
+from typing import Any, Optional
+
+from pydantic import BaseModel, Field
+
+
+class Distribution(BaseModel):
+    """Value distribution for a column."""
+
+    column: str = Field(..., description="Column name")
+    total_rows: int = Field(..., description="Total rows analyzed")
+    unique_values: int = Field(..., description="Number of unique values")
+    null_count: int = Field(..., description="Number of NULL values")
+    top_values: list[dict[str, Any]] = Field(
+        ..., description="Top N most frequent values with counts"
+    )
+    sample_size: int = Field(..., description="Number of rows sampled")
+
+    @property
+    def null_percentage(self) -> float:
+        """Percentage of NULL values."""
+        if self.total_rows == 0:
+            return 0.0
+        return (self.null_count / self.total_rows) * 100
+
+    @property
+    def cardinality(self) -> float:
+        """Cardinality (unique values / total rows)."""
+        if self.total_rows == 0:
+            return 0.0
+        return self.unique_values / self.total_rows
+
+    @property
+    def is_high_cardinality(self) -> bool:
+        """Check if column has high cardinality (>0.9)."""
+        return self.cardinality > 0.9
+
+    @property
+    def is_low_cardinality(self) -> bool:
+        """Check if column has low cardinality (<0.1)."""
+        return self.cardinality < 0.1
+
+
+class ColumnStats(BaseModel):
+    """Statistical information about a column."""
+
+    column: str = Field(..., description="Column name")
+    data_type: str = Field(..., description="Column data type")
+    total_rows: int = Field(..., description="Total rows in table")
+    null_count: int = Field(..., description="Number of NULL values")
+    distinct_count: Optional[int] = Field(None, description="Number of distinct values")
+    min_value: Optional[Any] = Field(None, description="Minimum value")
+    max_value: Optional[Any] = Field(None, description="Maximum value")
+    avg_value: Optional[float] = Field(
+        None, description="Average value (numeric columns)"
+    )
+    median_value: Optional[Any] = Field(None, description="Median value")
+    stddev_value: Optional[float] = Field(None, description="Standard deviation")
+    percentile_25: Optional[Any] = Field(None, description="25th percentile")
+    percentile_75: Optional[Any] = Field(None, description="75th percentile")
+    percentile_95: Optional[Any] = Field(None, description="95th percentile")
+    percentile_99: Optional[Any] = Field(None, description="99th percentile")
+    most_common_values: list[dict[str, Any]] = Field(
+        default_factory=list, description="Most common values with frequencies"
+    )
+    sample_size: int = Field(..., description="Number of rows sampled for statistics")
+    warning: Optional[str] = Field(
+        None, description="Warning message if stats unavailable"
+    )
+
+    @property
+    def null_percentage(self) -> float:
+        """Percentage of NULL values."""
+        if self.total_rows == 0:
+            return 0.0
+        return (self.null_count / self.total_rows) * 100
+
+    @property
+    def completeness(self) -> float:
+        """Data completeness (1 - null percentage as decimal)."""
+        return 1.0 - (self.null_percentage / 100.0)
+
+    @property
+    def cardinality(self) -> Optional[float]:
+        """Cardinality ratio (distinct / total)."""
+        if self.distinct_count is None or self.total_rows == 0:
+            return None
+        return self.distinct_count / self.total_rows
+
+    @property
+    def has_advanced_stats(self) -> bool:
+        """Check if advanced statistics are available."""
+        return any(
+            [
+                self.median_value is not None,
+                self.percentile_25 is not None,
+                self.stddev_value is not None,
+            ]
+        )
+
+    @property
+    def is_numeric(self) -> bool:
+        """Check if column appears to be numeric based on available stats."""
+        return self.avg_value is not None or self.stddev_value is not None
+
+    @property
+    def range_value(self) -> Optional[Any]:
+        """Calculate range (max - min) for numeric columns."""
+        if self.min_value is not None and self.max_value is not None:
+            try:
+                return self.max_value - self.min_value
+            except (TypeError, ValueError):
+                return None
+        return None
+
+    def get_quality_score(self) -> float:
+        """
+        Calculate data quality score (0-100).
+        Based on completeness, cardinality, and availability of stats.
+        """
+        score = 0.0
+
+        # Completeness (0-40 points)
+        score += self.completeness * 40
+
+        # Has distinct count (10 points)
+        if self.distinct_count is not None:
+            score += 10
+
+        # Has min/max (10 points)
+        if self.min_value is not None and self.max_value is not None:
+            score += 10
+
+        # Has advanced stats (20 points)
+        if self.has_advanced_stats:
+            score += 20
+
+        # Has most common values (10 points)
+        if self.most_common_values:
+            score += 10
+
+        # Has average (10 points if numeric)
+        if self.avg_value is not None:
+            score += 10
+
+        return min(score, 100.0)
+
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "column": "age",
+                    "data_type": "integer",
+                    "total_rows": 10000,
+                    "null_count": 50,
+                    "distinct_count": 95,
+                    "min_value": 18,
+                    "max_value": 95,
+                    "avg_value": 42.5,
+                    "median_value": 41,
+                    "stddev_value": 15.2,
+                    "percentile_25": 30,
+                    "percentile_75": 55,
+                    "percentile_95": 72,
+                    "percentile_99": 85,
+                    "most_common_values": [
+                        {"value": 35, "count": 250},
+                        {"value": 42, "count": 230},
+                    ],
+                    "sample_size": 10000,
+                    "warning": None,
+                }
+            ]
+        }
+    }

db_connect_mcp/models/table.py

@@ -0,0 +1,230 @@
+"""Table, column, index, and constraint information models."""
+
+import warnings
+from typing import Any, Optional
+
+from pydantic import BaseModel, ConfigDict, Field
+
+# Suppress the specific warning about field 'schema' shadowing
+warnings.filterwarnings(
+    "ignore",
+    message='Field name "schema" in "TableInfo" shadows an attribute in parent',
+    category=UserWarning,
+)
+
+
+class ColumnInfo(BaseModel):
+    """Information about a table column."""
+
+    name: str = Field(..., description="Column name")
+    data_type: str = Field(..., description="Column data type")
+    nullable: bool = Field(..., description="Whether column allows NULL")
+    default: Optional[str] = Field(None, description="Default value expression")
+    primary_key: bool = Field(
+        default=False, description="Whether column is part of primary key"
+    )
+    foreign_key: Optional[str] = Field(
+        None, description="Foreign key reference (table.column)"
+    )
+    unique: bool = Field(
+        default=False, description="Whether column has UNIQUE constraint"
+    )
+    indexed: bool = Field(default=False, description="Whether column is indexed")
+    comment: Optional[str] = Field(None, description="Column comment/description")
+    max_length: Optional[int] = Field(
+        None, description="Maximum length for string types"
+    )
+    numeric_precision: Optional[int] = Field(
+        None, description="Precision for numeric types"
+    )
+    numeric_scale: Optional[int] = Field(None, description="Scale for numeric types")
+    extra_info: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Database-specific additional information",
+    )
+
+
+class IndexInfo(BaseModel):
+    """Information about a table index."""
+
+    name: str = Field(..., description="Index name")
+    columns: list[str] = Field(..., description="Indexed column names")
+    unique: bool = Field(default=False, description="Whether index enforces uniqueness")
+    primary: bool = Field(
+        default=False, description="Whether this is the primary key index"
+    )
+    index_type: Optional[str] = Field(
+        None, description="Index type (btree, hash, etc.)"
+    )
+    size_bytes: Optional[int] = Field(None, description="Index size in bytes")
+    comment: Optional[str] = Field(None, description="Index comment")
+    extra_info: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Database-specific additional information",
+    )
+
+    @property
+    def size_human(self) -> Optional[str]:
+        """Human-readable size."""
+        if self.size_bytes is None:
+            return None
+
+        size = float(self.size_bytes)
+        for unit in ["B", "KB", "MB", "GB"]:
+            if size < 1024.0:
+                return f"{size:.2f} {unit}"
+            size /= 1024.0
+        return f"{size:.2f} TB"
+
+
+class ConstraintInfo(BaseModel):
+    """Information about a table constraint."""
+
+    name: str = Field(..., description="Constraint name")
+    constraint_type: str = Field(
+        ..., description="Constraint type (PRIMARY KEY, FOREIGN KEY, UNIQUE, CHECK)"
+    )
+    columns: list[str] = Field(..., description="Constrained column names")
+    referenced_table: Optional[str] = Field(None, description="Referenced table for FK")
+    referenced_columns: Optional[list[str]] = Field(
+        None, description="Referenced columns for FK"
+    )
+    definition: Optional[str] = Field(None, description="Constraint definition SQL")
+    deferrable: Optional[bool] = Field(
+        None, description="Whether constraint is deferrable"
+    )
+    initially_deferred: Optional[bool] = Field(
+        None, description="Whether initially deferred"
+    )
+    extra_info: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Database-specific additional information",
+    )
+
+
+class RelationshipInfo(BaseModel):
+    """Information about a foreign key relationship between tables."""
+
+    from_table: str = Field(..., description="Source table name")
+    from_schema: Optional[str] = Field(None, description="Source schema name")
+    from_columns: list[str] = Field(..., description="Source column names")
+    to_table: str = Field(..., description="Target table name")
+    to_schema: Optional[str] = Field(None, description="Target schema name")
+    to_columns: list[str] = Field(..., description="Target column names")
+    constraint_name: str = Field(..., description="Foreign key constraint name")
+    on_delete: Optional[str] = Field(None, description="ON DELETE action")
+    on_update: Optional[str] = Field(None, description="ON UPDATE action")
+
+
+class TableInfo(BaseModel):
+    """Comprehensive information about a table."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    name: str = Field(..., description="Table name")
+    schema: Optional[str] = Field(default=None, description="Schema name")
+    table_type: str = Field(
+        default="BASE TABLE", description="Type (BASE TABLE, VIEW, etc.)"
+    )
+    row_count: Optional[int] = Field(None, description="Approximate row count")
+    size_bytes: Optional[int] = Field(None, description="Table size in bytes")
+    index_size_bytes: Optional[int] = Field(
+        None, description="Total index size in bytes"
+    )
+    columns: list[ColumnInfo] = Field(
+        default_factory=list, description="Column information"
+    )
+    indexes: list[IndexInfo] = Field(
+        default_factory=list, description="Index information"
+    )
+    constraints: list[ConstraintInfo] = Field(
+        default_factory=list, description="Constraint information"
+    )
+    comment: Optional[str] = Field(None, description="Table comment/description")
+    created_at: Optional[str] = Field(None, description="Creation timestamp")
+    updated_at: Optional[str] = Field(None, description="Last update timestamp")
+    owner: Optional[str] = Field(None, description="Table owner")
+    extra_info: dict[str, Any] = Field(
+        default_factory=dict,
+        description="Database-specific additional information (engine, partitions, etc.)",
+    )
+
+    @property
+    def size_human(self) -> Optional[str]:
+        """Human-readable table size."""
+        if self.size_bytes is None:
+            return None
+
+        size = float(self.size_bytes)
+        for unit in ["B", "KB", "MB", "GB", "TB"]:
+            if size < 1024.0:
+                return f"{size:.2f} {unit}"
+            size /= 1024.0
+        return f"{size:.2f} PB"
+
+    @property
+    def total_size_bytes(self) -> Optional[int]:
+        """Total size including indexes."""
+        if self.size_bytes is None:
+            return None
+        return self.size_bytes + (self.index_size_bytes or 0)
+
+    @property
+    def total_size_human(self) -> Optional[str]:
+        """Human-readable total size."""
+        total = self.total_size_bytes
+        if total is None:
+            return None
+
+        size = float(total)
+        for unit in ["B", "KB", "MB", "GB", "TB"]:
+            if size < 1024.0:
+                return f"{size:.2f} {unit}"
+            size /= 1024.0
+        return f"{size:.2f} PB"
+
+    @property
+    def primary_key_columns(self) -> list[str]:
+        """Get primary key column names."""
+        return [col.name for col in self.columns if col.primary_key]
+
+    @property
+    def foreign_key_columns(self) -> list[str]:
+        """Get foreign key column names."""
+        return [col.name for col in self.columns if col.foreign_key is not None]
+
+    @property
+    def column_count(self) -> int:
+        """Get number of columns."""
+        return len(self.columns)
+
+    @property
+    def index_count(self) -> int:
+        """Get number of indexes."""
+        return len(self.indexes)
+
+    @property
+    def constraint_count(self) -> int:
+        """Get number of constraints."""
+        return len(self.constraints)
+
+    def get_column(self, name: str) -> Optional[ColumnInfo]:
+        """Get column by name."""
+        for col in self.columns:
+            if col.name == name:
+                return col
+        return None
+
+    def get_index(self, name: str) -> Optional[IndexInfo]:
+        """Get index by name."""
+        for idx in self.indexes:
+            if idx.name == name:
+                return idx
+        return None
+
+    def get_constraint(self, name: str) -> Optional[ConstraintInfo]:
+        """Get constraint by name."""
+        for constraint in self.constraints:
+            if constraint.name == name:
+                return constraint
+        return None