feat: Initial migration for Workspaces and pipeline step

migrations/versions/5c2f3eee5f90_introduce_workspaces.py

@@ -0,0 +1,61 @@
+"""introduce workspaces
+
+Revision ID: 5c2f3eee5f90
+Revises: 30d0144e1a50
+Create Date: 2025-01-15 19:27:08.230296
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "5c2f3eee5f90"
+down_revision: Union[str, None] = "30d0144e1a50"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+ # Workspaces table
+ op.execute(
+ """
+ CREATE TABLE workspaces (
+ id TEXT PRIMARY KEY, -- UUID stored as TEXT
+ name TEXT NOT NULL,
+ UNIQUE (name)
+ );
+ """
+ )
+ op.execute("INSERT INTO workspaces (id, name) VALUES ('1', 'default');")
+ # Sessions table
+ op.execute(
+ """
+ CREATE TABLE sessions (
+ id TEXT PRIMARY KEY, -- UUID stored as TEXT
+ active_workspace_id TEXT NOT NULL,
+ last_update DATETIME NOT NULL,
+ FOREIGN KEY (active_workspace_id) REFERENCES workspaces(id)
+ );
+ """
+ )
+ # Alter table prompts
+ op.execute("ALTER TABLE prompts ADD COLUMN workspace_id TEXT REFERENCES workspaces(id);")
+ op.execute("UPDATE prompts SET workspace_id = '1';")
+ # Create index for workspace_id
+ op.execute("CREATE INDEX idx_prompts_workspace_id ON prompts (workspace_id);")
+ # Create index for session_id
+ op.execute("CREATE INDEX idx_sessions_workspace_id ON sessions (active_workspace_id);")
+
+
+def downgrade() -> None:
+ # Drop the index for workspace_id
+ op.execute("DROP INDEX IF EXISTS idx_prompts_workspace_id;")
+ op.execute("DROP INDEX IF EXISTS idx_sessions_workspace_id;")
+ # Remove the workspace_id column from prompts table
+ op.execute("ALTER TABLE prompts DROP COLUMN workspace_id;")
+ # Drop the sessions table
+ op.execute("DROP TABLE IF EXISTS sessions;")
+ # Drop the workspaces table
+ op.execute("DROP TABLE IF EXISTS workspaces;")

src/codegate/__init__.py

@@ -10,6 +10,7 @@
 _VERSION = "dev"
 _DESC = "CodeGate - A Generative AI security gateway."
 
+
 def __get_version_and_description() -> tuple[str, str]:
  try:
  version = metadata.version("codegate")
@@ -19,6 +20,7 @@ def __get_version_and_description() -> tuple[str, str]:
  description = _DESC
  return version, description
 
+
 __version__, __description__ = __get_version_and_description()
 
 __all__ = ["Config", "ConfigurationError", "LogFormat", "LogLevel", "setup_logging"]

src/codegate/cli.py

@@ -14,7 +14,7 @@
 from codegate.ca.codegate_ca import CertificateAuthority
 from codegate.codegate_logging import LogFormat, LogLevel, setup_logging
 from codegate.config import Config, ConfigurationError
-from codegate.db.connection import init_db_sync
+from codegate.db.connection import init_db_sync, init_session_if_not_exists
 from codegate.pipeline.factory import PipelineFactory
 from codegate.pipeline.secrets.manager import SecretsManager
 from codegate.providers.copilot.provider import CopilotProvider
@@ -307,6 +307,7 @@ def serve(
  logger = structlog.get_logger("codegate").bind(origin="cli")
 
  init_db_sync(cfg.db_path)
+ init_session_if_not_exists(cfg.db_path)
 
  # Check certificates and create CA if necessary
  logger.info("Checking certificates and creating CA if needed")

src/codegate/db/connection.py

@@ -1,23 +1,28 @@
 import asyncio
 import json
+import uuid
 from pathlib import Path
 from typing import List, Optional, Type
 
 import structlog
 from alembic import command as alembic_command
 from alembic.config import Config as AlembicConfig
-from pydantic import BaseModel
-from sqlalchemy import TextClause, text
+from pydantic import BaseModel, ValidationError
+from sqlalchemy import CursorResult, TextClause, text
 from sqlalchemy.exc import OperationalError
 from sqlalchemy.ext.asyncio import create_async_engine
 
 from codegate.db.fim_cache import FimCache
 from codegate.db.models import (
+ ActiveWorkspace,
  Alert,
  GetAlertsWithPromptAndOutputRow,
  GetPromptWithOutputsRow,
  Output,
  Prompt,
+ Session,
+ Workspace,
+ WorkspaceActive,
 )
 from codegate.pipeline.base import PipelineContext
 
@@ -75,10 +80,14 @@ async def _execute_update_pydantic_model(
  async def record_request(self, prompt_params: Optional[Prompt] = None) -> Optional[Prompt]:
  if prompt_params is None:
  return None
+ # Get the active workspace to store the request
+ active_workspace = await DbReader().get_active_workspace()
+ workspace_id = active_workspace.id if active_workspace else "1"
+ prompt_params.workspace_id = workspace_id
  sql = text(
  """
- INSERT INTO prompts (id, timestamp, provider, request, type)
- VALUES (:id, :timestamp, :provider, :request, :type)
+ INSERT INTO prompts (id, timestamp, provider, request, type, workspace_id)
+ VALUES (:id, :timestamp, :provider, :request, :type, :workspace_id)
  RETURNING *
  """
  )
@@ -223,26 +232,78 @@ async def record_context(self, context: Optional[PipelineContext]) -> None:
  except Exception as e:
  logger.error(f"Failed to record context: {context}.", error=str(e))
 
+ async def add_workspace(self, workspace_name: str) -> Optional[Workspace]:
+ try:
+ workspace = Workspace(id=str(uuid.uuid4()), name=workspace_name)
+ except ValidationError as e:
+ logger.error(f"Failed to create workspace with name: {workspace_name}: {str(e)}")
+ return None
+
+ sql = text(
+ """
+ INSERT INTO workspaces (id, name)
+ VALUES (:id, :name)
+ RETURNING *
+ """
+ )
+ added_workspace = await self._execute_update_pydantic_model(workspace, sql)
+ return added_workspace
+
+ async def update_session(self, session: Session) -> Optional[Session]:
+ sql = text(
+ """
+ INSERT INTO sessions (id, active_workspace_id, last_update)
+ VALUES (:id, :active_workspace_id, :last_update)
+ ON CONFLICT (id) DO UPDATE SET
+ active_workspace_id = excluded.active_workspace_id, last_update = excluded.last_update
+ WHERE id = excluded.id
+ RETURNING *
+ """
+ )
+ # We only pass an object to respect the signature of the function
+ active_session = await self._execute_update_pydantic_model(session, sql)
+ return active_session
+
 
 class DbReader(DbCodeGate):
 
  def __init__(self, sqlite_path: Optional[str] = None):
  super().__init__(sqlite_path)
 
+ async def _dump_result_to_pydantic_model(
+ self, model_type: Type[BaseModel], result: CursorResult
+ ) -> Optional[List[BaseModel]]:
+ try:
+ if not result:
+ return None
+ rows = [model_type(**row._asdict()) for row in result.fetchall() if row]
+ return rows
+ except Exception as e:
+ logger.error(f"Failed to dump to pydantic model: {model_type}.", error=str(e))
+ return None
+
  async def _execute_select_pydantic_model(
  self, model_type: Type[BaseModel], sql_command: TextClause
- ) -> Optional[BaseModel]:
+ ) -> Optional[List[BaseModel]]:
  async with self._async_db_engine.begin() as conn:
  try:
  result = await conn.execute(sql_command)
- if not result:
- return None
- rows = [model_type(**row._asdict()) for row in result.fetchall() if row]
- return rows
+ return await self._dump_result_to_pydantic_model(model_type, result)
  except Exception as e:
  logger.error(f"Failed to select model: {model_type}.", error=str(e))
  return None
 
+ async def _exec_select_conditions_to_pydantic(
+ self, model_type: Type[BaseModel], sql_command: TextClause, conditions: dict
+ ) -> Optional[List[BaseModel]]:
+ async with self._async_db_engine.begin() as conn:
+ try:
+ result = await conn.execute(sql_command, conditions)
+ return await self._dump_result_to_pydantic_model(model_type, result)
+ except Exception as e:
+ logger.error(f"Failed to select model with conditions: {model_type}.", error=str(e))
+ return None
+
  async def get_prompts_with_output(self) -> List[GetPromptWithOutputsRow]:
  sql = text(
  """
@@ -286,6 +347,54 @@ async def get_alerts_with_prompt_and_output(self) -> List[GetAlertsWithPromptAnd
  prompts = await self._execute_select_pydantic_model(GetAlertsWithPromptAndOutputRow, sql)
  return prompts
 
+ async def get_workspaces(self) -> List[WorkspaceActive]:
+ sql = text(
+ """
+ SELECT
+ w.id, w.name, s.active_workspace_id
+ FROM workspaces w
+ LEFT JOIN sessions s ON w.id = s.active_workspace_id
+ """
+ )
+ workspaces = await self._execute_select_pydantic_model(WorkspaceActive, sql)
+ return workspaces
+
+ async def get_workspace_by_name(self, name: str) -> List[Workspace]:
+ sql = text(
+ """
+ SELECT
+ id, name
+ FROM workspaces
+ WHERE name = :name
+ """
+ )
+ conditions = {"name": name}
+ workspaces = await self._exec_select_conditions_to_pydantic(Workspace, sql, conditions)
+ return workspaces[0] if workspaces else None
+
+ async def get_sessions(self) -> List[Session]:
+ sql = text(
+ """
+ SELECT
+ id, active_workspace_id, last_update
+ FROM sessions
+ """
+ )
+ sessions = await self._execute_select_pydantic_model(Session, sql)
+ return sessions
+
+ async def get_active_workspace(self) -> Optional[ActiveWorkspace]:
+ sql = text(
+ """
+ SELECT
+ w.id, w.name, s.id as session_id, s.last_update
+ FROM sessions s
+ INNER JOIN workspaces w ON w.id = s.active_workspace_id
+ """
+ )
+ active_workspace = await self._execute_select_pydantic_model(ActiveWorkspace, sql)
+ return active_workspace[0] if active_workspace else None
+
 
 def init_db_sync(db_path: Optional[str] = None):
  """DB will be initialized in the constructor in case it doesn't exist."""
@@ -307,5 +416,23 @@ def init_db_sync(db_path: Optional[str] = None):
  logger.info("DB initialized successfully.")
 
 
+def init_session_if_not_exists(db_path: Optional[str] = None):
+ import datetime
+
+ db_reader = DbReader(db_path)
+ sessions = asyncio.run(db_reader.get_sessions())
+ # If there are no sessions, create a new one
+ # TODO: For the moment there's a single session. If it already exists, we don't create a new one
+ if not sessions:
+ session = Session(
+ id=str(uuid.uuid4()),
+ active_workspace_id="1",
+ last_update=datetime.datetime.now(datetime.timezone.utc),
+ )
+ db_recorder = DbRecorder(db_path)
+ asyncio.run(db_recorder.update_session(session))
+ logger.info("Session in DB initialized successfully.")
+
+
 if __name__ == "__main__":
  init_db_sync()