prismiq 0.1.0__py3-none-any.whl

prismiq/schema.py

@@ -0,0 +1,333 @@
+"""Schema introspection for PostgreSQL databases.
+
+This module provides the SchemaIntrospector class that reads database
+metadata from PostgreSQL's information_schema.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from prismiq.types import (
+    ColumnSchema,
+    DatabaseSchema,
+    Relationship,
+    TableNotFoundError,
+    TableSchema,
+)
+
+if TYPE_CHECKING:
+    from asyncpg import Pool, Record  # type: ignore[import-not-found]
+
+    from prismiq.cache import CacheBackend
+
+
+class SchemaIntrospector:
+    """Introspects PostgreSQL database schema.
+
+    Reads table and column metadata from information_schema,
+    detects foreign key relationships, and provides a filtered
+    view based on exposed_tables configuration.
+
+    Supports optional caching to reduce database queries.
+
+    Example:
+        >>> pool = await asyncpg.create_pool(database_url)
+        >>> introspector = SchemaIntrospector(pool, exposed_tables=["users", "orders"])
+        >>> schema = await introspector.get_schema()
+        >>> print(schema.table_names())
+        ['users', 'orders']
+
+    With caching:
+        >>> from prismiq.cache import InMemoryCache
+        >>> cache = InMemoryCache()
+        >>> introspector = SchemaIntrospector(pool, cache=cache, cache_ttl=3600)
+        >>> schema = await introspector.get_schema()  # Hits database
+        >>> schema = await introspector.get_schema()  # Returns cached result
+    """
+
+    def __init__(
+        self,
+        pool: Pool,
+        exposed_tables: list[str] | None = None,
+        schema_name: str = "public",
+        cache: CacheBackend | None = None,
+        cache_ttl: int = 3600,
+    ) -> None:
+        """Initialize the schema introspector.
+
+        Args:
+            pool: asyncpg connection pool to use for queries.
+            exposed_tables: List of table names to expose. If None, all tables
+                in the schema are exposed.
+            schema_name: PostgreSQL schema to introspect (default: "public").
+            cache: Optional cache backend for caching schema data.
+            cache_ttl: TTL for cached schema in seconds (default: 1 hour).
+        """
+        self._pool = pool
+        self._exposed_tables = exposed_tables
+        self._schema_name = schema_name
+        self._cache = cache
+        self._cache_ttl = cache_ttl
+
+    def _cache_key(self, suffix: str) -> str:
+        """Generate schema-qualified cache key for tenant isolation.
+
+        Args:
+            suffix: Cache key suffix (e.g., "full", "table:users").
+
+        Returns:
+            Cache key with schema prefix (e.g., "schema:org_123:full").
+        """
+        return f"schema:{self._schema_name}:{suffix}"
+
+    async def get_schema(self, force_refresh: bool = False) -> DatabaseSchema:
+        """Get the complete database schema.
+
+        Args:
+            force_refresh: If True, bypass cache and fetch fresh data.
+
+        Returns:
+            DatabaseSchema containing all exposed tables and their relationships.
+        """
+        # Try cache first (using schema-qualified key for tenant isolation)
+        if self._cache and not force_refresh:
+            cached = await self._cache.get(self._cache_key("full"))
+            if cached is not None:
+                return DatabaseSchema.model_validate(cached)
+
+        # Introspect from database
+        schema = await self._introspect_schema()
+
+        # Store in cache (using schema-qualified key)
+        if self._cache:
+            await self._cache.set(self._cache_key("full"), schema.model_dump(), self._cache_ttl)
+
+        return schema
+
+    async def _introspect_schema(self) -> DatabaseSchema:
+        """Introspect schema from database."""
+        table_names = await self._get_table_names()
+        tables: list[TableSchema] = []
+
+        for table_name in table_names:
+            table = await self._get_table_schema(table_name)
+            tables.append(table)
+
+        relationships = await self.detect_relationships()
+
+        return DatabaseSchema(tables=tables, relationships=relationships)
+
+    async def get_table(self, table_name: str, force_refresh: bool = False) -> TableSchema:
+        """Get schema information for a single table.
+
+        Args:
+            table_name: Name of the table to retrieve.
+            force_refresh: If True, bypass cache and fetch fresh data.
+
+        Returns:
+            TableSchema for the requested table.
+
+        Raises:
+            TableNotFoundError: If the table doesn't exist or isn't exposed.
+        """
+        # Check if table is exposed
+        if self._exposed_tables is not None and table_name not in self._exposed_tables:
+            raise TableNotFoundError(table_name)
+
+        # Try cache first (using schema-qualified key for tenant isolation)
+        cache_key = self._cache_key(f"table:{table_name}")
+        if self._cache and not force_refresh:
+            cached = await self._cache.get(cache_key)
+            if cached is not None:
+                return TableSchema.model_validate(cached)
+
+        # Check if table exists in database
+        table_names = await self._get_table_names()
+        if table_name not in table_names:
+            raise TableNotFoundError(table_name)
+
+        # Introspect from database
+        table = await self._get_table_schema(table_name)
+
+        # Store in cache (using schema-qualified key)
+        if self._cache:
+            await self._cache.set(cache_key, table.model_dump(), self._cache_ttl)
+
+        return table
+
+    async def invalidate_cache(self) -> int:
+        """Invalidate all cached schema data for this schema.
+
+        Only invalidates cache entries for this schema (tenant isolation).
+
+        Returns:
+            Number of cache entries cleared.
+        """
+        if self._cache is None:
+            return 0
+
+        # Only clear cache entries for this specific schema
+        return await self._cache.clear(f"schema:{self._schema_name}:*")
+
+    async def detect_relationships(self) -> list[Relationship]:
+        """Detect foreign key relationships between exposed tables.
+
+        Returns:
+            List of Relationship objects representing foreign keys.
+        """
+        async with self._pool.acquire() as conn:
+            # Query foreign key constraints
+            query = """
+                SELECT
+                    tc.table_name AS from_table,
+                    kcu.column_name AS from_column,
+                    ccu.table_name AS to_table,
+                    ccu.column_name AS to_column
+                FROM information_schema.table_constraints tc
+                JOIN information_schema.key_column_usage kcu
+                    ON tc.constraint_name = kcu.constraint_name
+                    AND tc.table_schema = kcu.table_schema
+                JOIN information_schema.constraint_column_usage ccu
+                    ON ccu.constraint_name = tc.constraint_name
+                    AND ccu.table_schema = tc.table_schema
+                WHERE tc.constraint_type = 'FOREIGN KEY'
+                    AND tc.table_schema = $1
+            """
+            rows: list[Record] = await conn.fetch(query, self._schema_name)
+
+        relationships: list[Relationship] = []
+        exposed_set = set(self._exposed_tables) if self._exposed_tables else None
+
+        for row in rows:
+            from_table = row["from_table"]
+            to_table = row["to_table"]
+
+            # Filter to only include relationships between exposed tables
+            if exposed_set is not None and (
+                from_table not in exposed_set or to_table not in exposed_set
+            ):
+                continue
+
+            relationships.append(
+                Relationship(
+                    from_table=from_table,
+                    from_column=row["from_column"],
+                    to_table=to_table,
+                    to_column=row["to_column"],
+                )
+            )
+
+        return relationships
+
+    async def _get_table_names(self) -> list[str]:
+        """Get list of table names in the schema."""
+        async with self._pool.acquire() as conn:
+            query = """
+                SELECT table_name
+                FROM information_schema.tables
+                WHERE table_schema = $1
+                    AND table_type IN ('BASE TABLE', 'VIEW')
+                ORDER BY table_name
+            """
+            rows: list[Record] = await conn.fetch(query, self._schema_name)
+
+        table_names = [row["table_name"] for row in rows]
+
+        # Filter to exposed tables if specified
+        if self._exposed_tables is not None:
+            exposed_set = set(self._exposed_tables)
+            table_names = [t for t in table_names if t in exposed_set]
+
+        return table_names
+
+    async def _get_table_schema(self, table_name: str) -> TableSchema:
+        """Get schema for a single table."""
+        columns = await self._get_columns(table_name)
+        primary_keys = await self._get_primary_keys(table_name)
+        row_count = await self._get_row_count(table_name)
+        primary_key_set = set(primary_keys)
+
+        # Mark primary key columns
+        for col in columns:
+            if col.name in primary_key_set:
+                # Create new column with is_primary_key=True
+                # (Pydantic models are immutable by default in strict mode)
+                col_dict = col.model_dump()
+                col_dict["is_primary_key"] = True
+                columns[columns.index(col)] = ColumnSchema(**col_dict)
+
+        return TableSchema(
+            name=table_name,
+            schema_name=self._schema_name,
+            columns=columns,
+            row_count=row_count,
+        )
+
+    async def _get_columns(self, table_name: str) -> list[ColumnSchema]:
+        """Get column information for a table."""
+        async with self._pool.acquire() as conn:
+            query = """
+                SELECT
+                    column_name,
+                    data_type,
+                    is_nullable,
+                    column_default
+                FROM information_schema.columns
+                WHERE table_schema = $1
+                    AND table_name = $2
+                ORDER BY ordinal_position
+            """
+            rows: list[Record] = await conn.fetch(query, self._schema_name, table_name)
+
+        return [
+            ColumnSchema(
+                name=row["column_name"],
+                data_type=row["data_type"],
+                is_nullable=row["is_nullable"] == "YES",
+                default_value=row["column_default"],
+            )
+            for row in rows
+        ]
+
+    async def _get_primary_keys(self, table_name: str) -> list[str]:
+        """Get primary key column names for a table."""
+        async with self._pool.acquire() as conn:
+            query = """
+                SELECT kcu.column_name
+                FROM information_schema.table_constraints tc
+                JOIN information_schema.key_column_usage kcu
+                    ON tc.constraint_name = kcu.constraint_name
+                    AND tc.table_schema = kcu.table_schema
+                WHERE tc.constraint_type = 'PRIMARY KEY'
+                    AND tc.table_schema = $1
+                    AND tc.table_name = $2
+                ORDER BY kcu.ordinal_position
+            """
+            rows: list[Record] = await conn.fetch(query, self._schema_name, table_name)
+
+        return [row["column_name"] for row in rows]
+
+    async def _get_row_count(self, table_name: str) -> int | None:
+        """Get approximate row count for a table using pg_class.reltuples.
+
+        This is fast but may be slightly out of date. For exact counts,
+        VACUUM ANALYZE should be run periodically.
+        """
+        async with self._pool.acquire() as conn:
+            query = """
+                SELECT reltuples::bigint AS row_count
+                FROM pg_class c
+                JOIN pg_namespace n ON n.oid = c.relnamespace
+                WHERE n.nspname = $1
+                    AND c.relname = $2
+                    AND c.relkind = 'r'
+            """
+            row = await conn.fetchrow(query, self._schema_name, table_name)
+
+        if row is None:
+            return None
+
+        # reltuples can be -1 if never analyzed, treat as 0
+        count = row["row_count"]
+        return max(0, count) if count is not None else None

prismiq/schema_config.py

@@ -0,0 +1,354 @@
+"""Schema customization for Prismiq analytics.
+
+This module provides configuration models and a manager for customizing
+how database schema is presented to users, including friendly names,
+hidden columns, and formatting hints.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict
+
+from prismiq.types import DatabaseSchema
+
+
+class ColumnConfig(BaseModel):
+    """Configuration for a single column."""
+
+    model_config = ConfigDict(strict=True)
+
+    display_name: str | None = None
+    """Friendly name for UI display."""
+
+    description: str | None = None
+    """Tooltip/help text for the column."""
+
+    hidden: bool = False
+    """Whether to hide this column from the schema explorer."""
+
+    format: str | None = None
+    """Number format: plain, currency, percent, compact."""
+
+    date_format: str | None = None
+    """Date format string for date/timestamp columns."""
+
+
+class TableConfig(BaseModel):
+    """Configuration for a single table."""
+
+    model_config = ConfigDict(strict=True)
+
+    display_name: str | None = None
+    """Friendly name for UI display."""
+
+    description: str | None = None
+    """Tooltip/help text for the table."""
+
+    hidden: bool = False
+    """Whether to hide this table from the schema explorer."""
+
+    columns: dict[str, ColumnConfig] = {}
+    """Column-specific configurations."""
+
+
+class SchemaConfig(BaseModel):
+    """Complete schema customization configuration."""
+
+    model_config = ConfigDict(strict=True)
+
+    tables: dict[str, TableConfig] = {}
+    """Table-specific configurations."""
+
+    def get_table_config(self, table_name: str) -> TableConfig:
+        """Get config for a table, with defaults if not configured.
+
+        Args:
+            table_name: The table name to get configuration for.
+
+        Returns:
+            TableConfig for the table (may be default/empty).
+        """
+        return self.tables.get(table_name, TableConfig())
+
+    def get_column_config(self, table_name: str, column_name: str) -> ColumnConfig:
+        """Get config for a column, with defaults if not configured.
+
+        Args:
+            table_name: The table containing the column.
+            column_name: The column name to get configuration for.
+
+        Returns:
+            ColumnConfig for the column (may be default/empty).
+        """
+        table_config = self.get_table_config(table_name)
+        return table_config.columns.get(column_name, ColumnConfig())
+
+    def get_display_name(self, table_name: str, column_name: str | None = None) -> str:
+        """Get display name for table or column, falling back to actual name.
+
+        Args:
+            table_name: The table name.
+            column_name: Optional column name. If None, returns table display name.
+
+        Returns:
+            The display name, or the actual name if no display name is configured.
+        """
+        if column_name is None:
+            table_config = self.get_table_config(table_name)
+            return table_config.display_name or table_name
+        else:
+            column_config = self.get_column_config(table_name, column_name)
+            return column_config.display_name or column_name
+
+    def is_table_hidden(self, table_name: str) -> bool:
+        """Check if a table is hidden."""
+        return self.get_table_config(table_name).hidden
+
+    def is_column_hidden(self, table_name: str, column_name: str) -> bool:
+        """Check if a column is hidden."""
+        return self.get_column_config(table_name, column_name).hidden
+
+
+class EnhancedColumnSchema(BaseModel):
+    """Column schema with configuration-based enhancements."""
+
+    model_config = ConfigDict(strict=True)
+
+    name: str
+    """Column name in the database."""
+
+    data_type: str
+    """PostgreSQL data type."""
+
+    is_nullable: bool
+    """Whether the column allows NULL values."""
+
+    is_primary_key: bool = False
+    """Whether this column is part of the primary key."""
+
+    default_value: str | None = None
+    """Default value expression, if any."""
+
+    display_name: str | None = None
+    """Friendly display name from configuration."""
+
+    description: str | None = None
+    """Description from configuration."""
+
+    format: str | None = None
+    """Number format from configuration."""
+
+    date_format: str | None = None
+    """Date format from configuration."""
+
+
+class EnhancedTableSchema(BaseModel):
+    """Table schema with configuration-based enhancements."""
+
+    model_config = ConfigDict(strict=True)
+
+    name: str
+    """Table name in the database."""
+
+    schema_name: str = "public"
+    """Database schema (namespace)."""
+
+    columns: list[EnhancedColumnSchema]
+    """Enhanced column schemas."""
+
+    display_name: str | None = None
+    """Friendly display name from configuration."""
+
+    description: str | None = None
+    """Description from configuration."""
+
+    def get_column(self, column_name: str) -> EnhancedColumnSchema | None:
+        """Get a column by name, or None if not found."""
+        for col in self.columns:
+            if col.name == column_name:
+                return col
+        return None
+
+    def has_column(self, column_name: str) -> bool:
+        """Check if the table has a column with the given name."""
+        return self.get_column(column_name) is not None
+
+
+class EnhancedDatabaseSchema(BaseModel):
+    """Database schema with configuration-based enhancements."""
+
+    model_config = ConfigDict(strict=True)
+
+    tables: list[EnhancedTableSchema]
+    """Enhanced table schemas."""
+
+    relationships: list[Any]  # Using Any to avoid circular import
+    """Foreign key relationships."""
+
+    def get_table(self, table_name: str) -> EnhancedTableSchema | None:
+        """Get a table by name, or None if not found."""
+        for table in self.tables:
+            if table.name == table_name:
+                return table
+        return None
+
+    def has_table(self, table_name: str) -> bool:
+        """Check if the schema contains a table with the given name."""
+        return self.get_table(table_name) is not None
+
+    def table_names(self) -> list[str]:
+        """Get list of all table names."""
+        return [t.name for t in self.tables]
+
+
+class SchemaConfigManager:
+    """Manages schema configuration persistence and application.
+
+    Provides methods for updating configuration and applying it to
+    database schemas.
+    """
+
+    def __init__(self, config: SchemaConfig | None = None) -> None:
+        """Initialize the schema config manager.
+
+        Args:
+            config: Initial configuration. If None, an empty config is used.
+        """
+        self._config = config or SchemaConfig()
+
+    def get_config(self) -> SchemaConfig:
+        """Get current configuration.
+
+        Returns:
+            The current SchemaConfig.
+        """
+        return self._config
+
+    def update_table_config(self, table_name: str, config: TableConfig) -> None:
+        """Update configuration for a table.
+
+        Creates a new config with the updated table (immutable operation).
+
+        Args:
+            table_name: Name of the table to configure.
+            config: New configuration for the table.
+        """
+        new_tables = dict(self._config.tables)
+        new_tables[table_name] = config
+        self._config = SchemaConfig(tables=new_tables)
+
+    def update_column_config(self, table_name: str, column_name: str, config: ColumnConfig) -> None:
+        """Update configuration for a column.
+
+        Creates a new config with the updated column (immutable operation).
+
+        Args:
+            table_name: Name of the table containing the column.
+            column_name: Name of the column to configure.
+            config: New configuration for the column.
+        """
+        # Get existing table config or create new one
+        table_config = self._config.get_table_config(table_name)
+        new_columns = dict(table_config.columns)
+        new_columns[column_name] = config
+
+        # Create new table config with updated columns
+        new_table_config = TableConfig(
+            display_name=table_config.display_name,
+            description=table_config.description,
+            hidden=table_config.hidden,
+            columns=new_columns,
+        )
+
+        # Update via table config
+        self.update_table_config(table_name, new_table_config)
+
+    def apply_to_schema(self, schema: DatabaseSchema) -> EnhancedDatabaseSchema:
+        """Apply configuration to a schema.
+
+        Adds display names, descriptions, and formats from configuration.
+        Filters out hidden tables and columns.
+
+        Args:
+            schema: The database schema to enhance.
+
+        Returns:
+            Enhanced schema with configuration applied and hidden items removed.
+        """
+        enhanced_tables: list[EnhancedTableSchema] = []
+
+        for table in schema.tables:
+            # Skip hidden tables
+            if self._config.is_table_hidden(table.name):
+                continue
+
+            table_config = self._config.get_table_config(table.name)
+            enhanced_columns: list[EnhancedColumnSchema] = []
+
+            for column in table.columns:
+                # Skip hidden columns
+                if self._config.is_column_hidden(table.name, column.name):
+                    continue
+
+                column_config = self._config.get_column_config(table.name, column.name)
+                enhanced_columns.append(
+                    EnhancedColumnSchema(
+                        name=column.name,
+                        data_type=column.data_type,
+                        is_nullable=column.is_nullable,
+                        is_primary_key=column.is_primary_key,
+                        default_value=column.default_value,
+                        display_name=column_config.display_name,
+                        description=column_config.description,
+                        format=column_config.format,
+                        date_format=column_config.date_format,
+                    )
+                )
+
+            enhanced_tables.append(
+                EnhancedTableSchema(
+                    name=table.name,
+                    schema_name=table.schema_name,
+                    columns=enhanced_columns,
+                    display_name=table_config.display_name,
+                    description=table_config.description,
+                )
+            )
+
+        # Filter relationships to only include visible tables
+        visible_table_names = {t.name for t in enhanced_tables}
+        visible_relationships = [
+            rel
+            for rel in schema.relationships
+            if rel.from_table in visible_table_names and rel.to_table in visible_table_names
+        ]
+
+        return EnhancedDatabaseSchema(
+            tables=enhanced_tables,
+            relationships=visible_relationships,
+        )
+
+    def to_json(self) -> str:
+        """Serialize configuration to JSON.
+
+        Returns:
+            JSON string representation of the configuration.
+        """
+        return self._config.model_dump_json(indent=2)
+
+    @classmethod
+    def from_json(cls, json_str: str) -> SchemaConfigManager:
+        """Deserialize configuration from JSON.
+
+        Args:
+            json_str: JSON string to parse.
+
+        Returns:
+            New SchemaConfigManager with the parsed configuration.
+        """
+        data = json.loads(json_str)
+        config = SchemaConfig.model_validate(data)
+        return cls(config)