kailash 0.3.2__py3-none-any.whl → 0.4.1__py3-none-any.whl

kailash/{api/database.py → middleware/database/models.py}

@@ -1,17 +1,13 @@
 """
-Database models and storage backend for Kailash Workflow Studio.
+Enhanced Database Models for Kailash Middleware
 
-This module provides:
-- SQLAlchemy models for workflows, nodes, executions, and user data
-- Database initialization and migration support
-- Repository classes for data access
+Consolidates and enhances existing database models from api/database.py
+with middleware-specific features and event integration.
 """
 
 import uuid
-from contextlib import contextmanager
-from datetime import UTC, datetime
-from pathlib import Path
-from typing import Any
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional
 
 from sqlalchemy import (
     JSON,
@@ -25,16 +21,14 @@ from sqlalchemy import (
     Text,
     create_engine,
 )
-from sqlalchemy.engine import Engine
-from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import Session, relationship, sessionmaker
+from sqlalchemy.orm import declarative_base, relationship
 from sqlalchemy.sql import func
 
 Base = declarative_base()
 
 
-class Workflow(Base):
-    """Workflow database model"""
+class WorkflowModel(Base):
+    """Enhanced workflow model with middleware integration"""
 
     __tablename__ = "workflows"
 
@@ -45,26 +39,41 @@ class Workflow(Base):
     definition = Column(JSON, nullable=False)
     version = Column(Integer, default=1)
     is_published = Column(Boolean, default=False)
+
+    # Middleware enhancements
+    session_id = Column(String(36), nullable=True, index=True)  # Associated session
+    agent_ui_config = Column(JSON)  # UI-specific configuration
+    real_time_enabled = Column(Boolean, default=True)  # Enable real-time updates
+
+    # Metadata
     created_by = Column(String(255))
     created_at = Column(DateTime, default=func.now())
     updated_at = Column(DateTime, default=func.now(), onupdate=func.now())
 
     # Relationships
     executions = relationship(
-        "WorkflowExecution", back_populates="workflow", cascade="all, delete-orphan"
+        "WorkflowExecutionModel",
+        back_populates="workflow",
+        cascade="all, delete-orphan",
     )
     versions = relationship(
-        "WorkflowVersion", back_populates="workflow", cascade="all, delete-orphan"
+        "WorkflowVersionModel", back_populates="workflow", cascade="all, delete-orphan"
+    )
+    permissions = relationship(
+        "WorkflowPermissionModel",
+        back_populates="workflow",
+        cascade="all, delete-orphan",
     )
 
     __table_args__ = (
         Index("idx_workflow_tenant_created", "tenant_id", "created_at"),
         Index("idx_workflow_tenant_name", "tenant_id", "name"),
+        Index("idx_workflow_session", "session_id"),
     )
 
 
-class WorkflowVersion(Base):
-    """Workflow version history"""
+class WorkflowVersionModel(Base):
+    """Enhanced workflow version with change tracking"""
 
     __tablename__ = "workflow_versions"
 
@@ -73,19 +82,25 @@ class WorkflowVersion(Base):
     version = Column(Integer, nullable=False)
     definition = Column(JSON, nullable=False)
     change_message = Column(Text)
+
+    # Middleware enhancements
+    diff_data = Column(JSON)  # Structured diff from previous version
+    migration_notes = Column(Text)  # Notes about breaking changes
+
+    # Metadata
     created_by = Column(String(255))
     created_at = Column(DateTime, default=func.now())
 
     # Relationships
-    workflow = relationship("Workflow", back_populates="versions")
+    workflow = relationship("WorkflowModel", back_populates="versions")
 
     __table_args__ = (
         Index("idx_version_workflow_version", "workflow_id", "version", unique=True),
     )
 
 
-class CustomNode(Base):
-    """Custom node definitions created by users"""
+class CustomNodeModel(Base):
+    """Enhanced custom node with middleware features"""
 
     __tablename__ = "custom_nodes"
 
@@ -103,9 +118,14 @@ class CustomNode(Base):
     outputs = Column(JSON)  # Output port definitions
 
     # Implementation
-    implementation_type = Column(String(50))  # 'python', 'workflow', 'api'
+    implementation_type = Column(String(50))  # 'python', 'workflow', 'api', 'kailash'
     implementation = Column(JSON)  # Implementation details
 
+    # Middleware enhancements
+    middleware_config = Column(JSON)  # Middleware-specific configuration
+    real_time_capable = Column(Boolean, default=False)  # Supports real-time updates
+    event_types = Column(JSON)  # List of event types this node can emit
+
     # Metadata
     is_published = Column(Boolean, default=False)
     created_by = Column(String(255))
@@ -118,8 +138,8 @@ class CustomNode(Base):
     )
 
 
-class WorkflowExecution(Base):
-    """Workflow execution history"""
+class WorkflowExecutionModel(Base):
+    """Enhanced execution model with real-time tracking"""
 
     __tablename__ = "workflow_executions"
 
@@ -135,6 +155,12 @@ class WorkflowExecution(Base):
     result = Column(JSON)
     error = Column(Text)
 
+    # Middleware enhancements
+    session_id = Column(String(36), nullable=True, index=True)  # Associated session
+    real_time_events = Column(JSON)  # List of real-time events emitted
+    agent_ui_data = Column(JSON)  # UI-specific execution data
+    execution_metrics = Column(JSON)  # Performance and resource metrics
+
     # Performance metrics
     started_at = Column(DateTime)
     completed_at = Column(DateTime)
@@ -144,17 +170,18 @@ class WorkflowExecution(Base):
     node_executions = Column(JSON)  # Detailed per-node execution data
 
     # Relationships
-    workflow = relationship("Workflow", back_populates="executions")
+    workflow = relationship("WorkflowModel", back_populates="executions")
 
     __table_args__ = (
         Index("idx_execution_tenant_started", "tenant_id", "started_at"),
         Index("idx_execution_workflow_started", "workflow_id", "started_at"),
         Index("idx_execution_status", "status"),
+        Index("idx_execution_session", "session_id"),
     )
 
 
-class UserPreferences(Base):
-    """User preferences and settings"""
+class UserPreferencesModel(Base):
+    """Enhanced user preferences with middleware features"""
 
     __tablename__ = "user_preferences"
 
@@ -171,6 +198,11 @@ class UserPreferences(Base):
     favorite_nodes = Column(JSON)
     recent_workflows = Column(JSON)
 
+    # Middleware enhancements
+    real_time_preferences = Column(JSON)  # Real-time update preferences
+    agent_ui_settings = Column(JSON)  # Agent-UI specific settings
+    notification_settings = Column(JSON)  # Event notification preferences
+
     created_at = Column(DateTime, default=func.now())
     updated_at = Column(DateTime, default=func.now(), onupdate=func.now())
 
@@ -179,8 +211,8 @@ class UserPreferences(Base):
     )
 
 
-class WorkflowTemplate(Base):
-    """Pre-built workflow templates"""
+class WorkflowTemplateModel(Base):
+    """Enhanced workflow templates with middleware integration"""
 
     __tablename__ = "workflow_templates"
 
@@ -195,6 +227,11 @@ class WorkflowTemplate(Base):
     definition = Column(JSON, nullable=False)
     default_parameters = Column(JSON)
 
+    # Middleware enhancements
+    middleware_features = Column(JSON)  # Required middleware features
+    agent_ui_template = Column(JSON)  # Agent-UI template configuration
+    real_time_template = Column(JSON)  # Real-time update template
+
     # Metadata
     is_public = Column(Boolean, default=False)
     usage_count = Column(Integer, default=0)
@@ -208,8 +245,8 @@ class WorkflowTemplate(Base):
     )
 
 
-class WorkflowPermission(Base):
-    """Workflow-level permissions"""
+class WorkflowPermissionModel(Base):
+    """Enhanced workflow permissions with middleware access control"""
 
     __tablename__ = "workflow_permissions"
 
@@ -233,13 +270,18 @@ class WorkflowPermission(Base):
     # Conditions (JSON object for flexibility)
     conditions = Column(JSON)
 
+    # Middleware enhancements
+    session_restrictions = Column(JSON)  # Session-based restrictions
+    real_time_permissions = Column(JSON)  # Real-time access permissions
+    agent_ui_permissions = Column(JSON)  # Agent-UI specific permissions
+
     # Metadata
     created_by = Column(String(255))
     created_at = Column(DateTime, default=func.now())
     expires_at = Column(DateTime, nullable=True)
 
     # Relationships
-    workflow = relationship("Workflow", backref="permissions")
+    workflow = relationship("WorkflowModel", back_populates="permissions")
 
     __table_args__ = (
         Index("idx_workflow_perm_workflow", "workflow_id"),
@@ -248,8 +290,8 @@ class WorkflowPermission(Base):
     )
 
 
-class NodePermission(Base):
-    """Node-level permissions"""
+class NodePermissionModel(Base):
+    """Enhanced node permissions with middleware features"""
 
     __tablename__ = "node_permissions"
 
@@ -274,13 +316,17 @@ class NodePermission(Base):
     masked_fields = Column(JSON)  # Fields to mask in output
     redirect_node = Column(String(255))  # Alternative node if access denied
 
+    # Middleware enhancements
+    real_time_masking = Column(JSON)  # Real-time output masking rules
+    event_filtering = Column(JSON)  # Event filtering rules for this node
+
     # Metadata
     created_by = Column(String(255))
     created_at = Column(DateTime, default=func.now())
     expires_at = Column(DateTime, nullable=True)
 
     # Relationships
-    workflow = relationship("Workflow", backref="node_permissions")
+    workflow = relationship("WorkflowModel", backref="node_permissions")
 
     __table_args__ = (
         Index("idx_node_perm_workflow", "workflow_id"),
@@ -290,8 +336,8 @@ class NodePermission(Base):
     )
 
 
-class AccessLog(Base):
-    """Audit log for access attempts"""
+class AccessLogModel(Base):
+    """Enhanced access log with middleware event integration"""
 
     __tablename__ = "access_logs"
 
@@ -300,7 +346,7 @@ class AccessLog(Base):
     user_id = Column(String(36), nullable=False, index=True)
 
     # What was accessed?
-    resource_type = Column(String(50), nullable=False)  # workflow, node
+    resource_type = Column(String(50), nullable=False)  # workflow, node, session, api
     resource_id = Column(String(255), nullable=False)
     permission = Column(String(50), nullable=False)
 
@@ -313,17 +359,25 @@ class AccessLog(Base):
     user_agent = Column(String(255))
     session_id = Column(String(36))
 
+    # Middleware enhancements
+    event_id = Column(String(36))  # Associated middleware event
+    real_time_logged = Column(
+        Boolean, default=False
+    )  # Whether logged to real-time stream
+    agent_ui_context = Column(JSON)  # Agent-UI specific context
+
     # Timestamp
     timestamp = Column(DateTime, default=func.now(), index=True)
 
     __table_args__ = (
         Index("idx_access_log_user_time", "user_id", "timestamp"),
         Index("idx_access_log_resource", "resource_type", "resource_id", "timestamp"),
+        Index("idx_access_log_event", "event_id"),
     )
 
 
-class UserGroup(Base):
-    """User groups for permission management"""
+class UserGroupModel(Base):
+    """Enhanced user groups with middleware features"""
 
     __tablename__ = "user_groups"
 
@@ -335,6 +389,11 @@ class UserGroup(Base):
     # Group permissions (can be inherited)
     permissions = Column(JSON)
 
+    # Middleware enhancements
+    middleware_permissions = Column(JSON)  # Middleware-specific permissions
+    real_time_access = Column(JSON)  # Real-time feature access
+    agent_ui_access = Column(JSON)  # Agent-UI access configuration
+
     # Metadata
     created_by = Column(String(255))
     created_at = Column(DateTime, default=func.now())
@@ -342,14 +401,14 @@ class UserGroup(Base):
 
     # Relationships
     members = relationship(
-        "UserGroupMember", back_populates="group", cascade="all, delete-orphan"
+        "UserGroupMemberModel", back_populates="group", cascade="all, delete-orphan"
     )
 
     __table_args__ = (Index("idx_group_tenant_name", "tenant_id", "name", unique=True),)
 
 
-class UserGroupMember(Base):
-    """User membership in groups"""
+class UserGroupMemberModel(Base):
+    """Enhanced group membership with middleware tracking"""
 
     __tablename__ = "user_group_members"
 
@@ -362,8 +421,11 @@ class UserGroupMember(Base):
     joined_at = Column(DateTime, default=func.now())
     added_by = Column(String(255))
 
+    # Middleware enhancements
+    middleware_role_overrides = Column(JSON)  # Role overrides for middleware features
+
     # Relationships
-    group = relationship("UserGroup", back_populates="members")
+    group = relationship("UserGroupModel", back_populates="members")
 
     __table_args__ = (
         Index("idx_group_member", "group_id", "user_id", unique=True),
@@ -371,250 +433,80 @@ class UserGroupMember(Base):
     )
 
 
-# Repository classes for data access
-class WorkflowRepository:
-    """Repository for workflow operations"""
-
-    def __init__(self, session: Session):
-        self.session = session
-
-    def create(
-        self,
-        tenant_id: str,
-        name: str,
-        description: str,
-        definition: dict[str, Any],
-        created_by: str = None,
-    ) -> Workflow:
-        """Create a new workflow"""
-        workflow = Workflow(
-            tenant_id=tenant_id,
-            name=name,
-            description=description,
-            definition=definition,
-            created_by=created_by,
-        )
-        self.session.add(workflow)
-        self.session.commit()
-
-        # Create initial version
-        self.create_version(workflow.id, 1, definition, "Initial version", created_by)
-
-        return workflow
-
-    def update(
-        self, workflow_id: str, updates: dict[str, Any], updated_by: str = None
-    ) -> Workflow:
-        """Update a workflow"""
-        workflow = self.session.query(Workflow).filter_by(id=workflow_id).first()
-        if not workflow:
-            raise ValueError(f"Workflow {workflow_id} not found")
-
-        # Update fields
-        for key, value in updates.items():
-            if hasattr(workflow, key):
-                setattr(workflow, key, value)
-
-        # If definition changed, create new version
-        if "definition" in updates:
-            workflow.version += 1
-            self.create_version(
-                workflow_id,
-                workflow.version,
-                updates["definition"],
-                updates.get("change_message", "Updated workflow"),
-                updated_by,
-            )
-
-        self.session.commit()
-        return workflow
-
-    def create_version(
-        self,
-        workflow_id: str,
-        version: int,
-        definition: dict[str, Any],
-        change_message: str,
-        created_by: str = None,
-    ):
-        """Create a workflow version"""
-        version_record = WorkflowVersion(
-            workflow_id=workflow_id,
-            version=version,
-            definition=definition,
-            change_message=change_message,
-            created_by=created_by,
-        )
-        self.session.add(version_record)
-        self.session.commit()
-
-    def get(self, workflow_id: str) -> Workflow | None:
-        """Get a workflow by ID"""
-        return self.session.query(Workflow).filter_by(id=workflow_id).first()
-
-    def list(self, tenant_id: str, limit: int = 100, offset: int = 0) -> list[Workflow]:
-        """List workflows for a tenant"""
-        return (
-            self.session.query(Workflow)
-            .filter_by(tenant_id=tenant_id)
-            .order_by(Workflow.created_at.desc())
-            .limit(limit)
-            .offset(offset)
-            .all()
-        )
-
-    def delete(self, workflow_id: str):
-        """Delete a workflow"""
-        workflow = self.get(workflow_id)
-        if workflow:
-            self.session.delete(workflow)
-            self.session.commit()
-
-
-class CustomNodeRepository:
-    """Repository for custom node operations"""
-
-    def __init__(self, session: Session):
-        self.session = session
-
-    def create(self, tenant_id: str, node_data: dict[str, Any]) -> CustomNode:
-        """Create a custom node"""
-        node = CustomNode(tenant_id=tenant_id, **node_data)
-        self.session.add(node)
-        self.session.commit()
-        return node
-
-    def update(self, node_id: str, updates: dict[str, Any]) -> CustomNode:
-        """Update a custom node"""
-        node = self.session.query(CustomNode).filter_by(id=node_id).first()
-        if not node:
-            raise ValueError(f"Custom node {node_id} not found")
-
-        for key, value in updates.items():
-            if hasattr(node, key):
-                setattr(node, key, value)
-
-        self.session.commit()
-        return node
-
-    def list(self, tenant_id: str) -> list[CustomNode]:
-        """List custom nodes for a tenant"""
-        return (
-            self.session.query(CustomNode)
-            .filter_by(tenant_id=tenant_id)
-            .order_by(CustomNode.category, CustomNode.name)
-            .all()
-        )
-
-    def get(self, node_id: str) -> CustomNode | None:
-        """Get a custom node by ID"""
-        return self.session.query(CustomNode).filter_by(id=node_id).first()
-
-    def delete(self, node_id: str):
-        """Delete a custom node"""
-        node = self.get(node_id)
-        if node:
-            self.session.delete(node)
-            self.session.commit()
-
-
-class ExecutionRepository:
-    """Repository for execution operations"""
-
-    def __init__(self, session: Session):
-        self.session = session
-
-    def create(
-        self, workflow_id: str, tenant_id: str, parameters: dict[str, Any] = None
-    ) -> WorkflowExecution:
-        """Create an execution record"""
-        execution = WorkflowExecution(
-            workflow_id=workflow_id,
-            tenant_id=tenant_id,
-            status="pending",
-            parameters=parameters,
-            started_at=datetime.now(UTC),
-        )
-        self.session.add(execution)
-        self.session.commit()
-        return execution
-
-    def update_status(
-        self,
-        execution_id: str,
-        status: str,
-        result: dict[str, Any] = None,
-        error: str = None,
-    ):
-        """Update execution status"""
-        execution = (
-            self.session.query(WorkflowExecution).filter_by(id=execution_id).first()
-        )
-        if not execution:
-            raise ValueError(f"Execution {execution_id} not found")
-
-        execution.status = status
-        if result is not None:
-            execution.result = result
-        if error is not None:
-            execution.error = error
-
-        if status in ["completed", "failed"]:
-            execution.completed_at = datetime.now(UTC)
-            if execution.started_at:
-                execution.execution_time_ms = int(
-                    (execution.completed_at - execution.started_at).total_seconds()
-                    * 1000
-                )
-
-        self.session.commit()
-
-    def get(self, execution_id: str) -> WorkflowExecution | None:
-        """Get execution by ID"""
-        return self.session.query(WorkflowExecution).filter_by(id=execution_id).first()
-
-    def list_for_workflow(
-        self, workflow_id: str, limit: int = 50
-    ) -> list[WorkflowExecution]:
-        """List executions for a workflow"""
-        return (
-            self.session.query(WorkflowExecution)
-            .filter_by(workflow_id=workflow_id)
-            .order_by(WorkflowExecution.started_at.desc())
-            .limit(limit)
-            .all()
-        )
-
-
-# Database initialization
-def init_database(db_path: str = None) -> tuple[sessionmaker, Engine]:
-    """Initialize the database"""
-    if db_path is None:
-        db_path = Path.home() / ".kailash" / "studio.db"
-
-    db_path = Path(db_path)
-    db_path.parent.mkdir(parents=True, exist_ok=True)
-
-    engine = create_engine(f"sqlite:///{db_path}", echo=False)
-    Base.metadata.create_all(engine)
-
-    SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
-
-    return SessionLocal, engine
-
-
-# Context manager for database sessions
-
-
-@contextmanager
-def get_db_session(SessionLocal):
-    """Provide a transactional scope for database operations"""
-    session = SessionLocal()
-    try:
-        yield session
-        session.commit()
-    except Exception:
-        session.rollback()
-        raise
-    finally:
-        session.close()
+class MiddlewareSessionModel(Base):
+    """Middleware-specific session tracking"""
+
+    __tablename__ = "middleware_sessions"
+
+    id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
+    tenant_id = Column(String(36), nullable=False, index=True)
+    user_id = Column(String(36), nullable=False, index=True)
+
+    # Session details
+    session_type = Column(String(50), nullable=False)  # agent_ui, api, websocket, mcp
+    status = Column(String(20), default="active")  # active, inactive, expired
+
+    # Session data
+    session_metadata = Column(JSON)  # Session-specific metadata
+    preferences = Column(JSON)  # User preferences for this session
+    context = Column(JSON)  # Session context data
+
+    # Connection details
+    connection_id = Column(String(255))  # WebSocket/API connection ID
+    ip_address = Column(String(50))
+    user_agent = Column(String(255))
+
+    # Timestamps
+    created_at = Column(DateTime, default=func.now())
+    last_activity = Column(DateTime, default=func.now())
+    expires_at = Column(DateTime)
+
+    __table_args__ = (
+        Index("idx_session_tenant_user", "tenant_id", "user_id"),
+        Index("idx_session_status", "status"),
+        Index("idx_session_type", "session_type"),
+        Index("idx_session_activity", "last_activity"),
+    )
+
+
+class MiddlewareEventModel(Base):
+    """Middleware event logging for audit and replay"""
+
+    __tablename__ = "middleware_events"
+
+    id = Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
+    tenant_id = Column(String(36), nullable=False, index=True)
+
+    # Event details
+    event_type = Column(String(100), nullable=False, index=True)
+    event_source = Column(
+        String(50), nullable=False
+    )  # agent_ui, api_gateway, realtime, etc.
+
+    # Event data
+    data = Column(JSON, nullable=False)
+    context = Column(JSON)  # Additional context
+
+    # Associated entities
+    session_id = Column(String(36), nullable=True, index=True)
+    user_id = Column(String(36), nullable=True, index=True)
+    workflow_id = Column(String(36), nullable=True, index=True)
+    execution_id = Column(String(36), nullable=True, index=True)
+
+    # Event metadata
+    correlation_id = Column(String(36))  # For tracing related events
+    sequence_number = Column(Integer)  # For ordering within correlation
+
+    # Processing status
+    processed = Column(Boolean, default=False)
+    processed_at = Column(DateTime)
+    error_message = Column(Text)
+
+    # Timestamp
+    timestamp = Column(DateTime, default=func.now(), index=True)
+
+    __table_args__ = (
+        Index("idx_event_type_time", "event_type", "timestamp"),
+        Index("idx_event_correlation", "correlation_id", "sequence_number"),
+        Index("idx_event_processed", "processed"),
+    )