Skip to content
This repository was archived by the owner on Jun 5, 2025. It is now read-only.

Commit 6c8e407

Browse files
committed
Add instance table along with init code.
This change adds an `instance` table containing the minimum set of details about codegate. We might want to add more details in the future. The table must contain only a single record at any time. To guarantee that, a trigger is added that prevents inserts beyond the first record. Finally, code performing the initialization is added to the `serve` command.
1 parent da69ec0 commit 6c8e407

File tree

4 files changed

+118
-3
lines changed

4 files changed

+118
-3
lines changed
Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
"""add installation table
2+
3+
Revision ID: e4c05d7591a8
4+
Revises: 3ec2b4ab569c
5+
Create Date: 2025-03-05 21:26:19.034319+00:00
6+
7+
"""
8+
9+
from typing import Sequence, Union
10+
11+
from alembic import op
12+
import sqlalchemy as sa
13+
14+
15+
# revision identifiers, used by Alembic.
16+
revision: str = "e4c05d7591a8"
17+
down_revision: Union[str, None] = "3ec2b4ab569c"
18+
branch_labels: Union[str, Sequence[str], None] = None
19+
depends_on: Union[str, Sequence[str], None] = None
20+
21+
22+
def upgrade() -> None:
23+
op.execute("BEGIN TRANSACTION;")
24+
25+
op.execute(
26+
"""
27+
CREATE TABLE IF NOT EXISTS instance (
28+
id TEXT PRIMARY KEY, -- UUID stored as TEXT
29+
created_at DATETIME NOT NULL
30+
);
31+
"""
32+
)
33+
34+
op.execute(
35+
"""
36+
-- The following trigger prevents multiple insertions in the
37+
-- instance table. It is safe since the dimension of the table
38+
-- is fixed.
39+
40+
CREATE TRIGGER single_instance
41+
BEFORE INSERT ON instance
42+
WHEN (SELECT COUNT(*) FROM instance) >= 1
43+
BEGIN
44+
SELECT RAISE(FAIL, 'only one instance!');
45+
END;
46+
"""
47+
)
48+
49+
# Finish transaction
50+
op.execute("COMMIT;")
51+
52+
53+
def downgrade() -> None:
54+
op.execute("BEGIN TRANSACTION;")
55+
56+
op.execute(
57+
"""
58+
DROP TABLE instance;
59+
"""
60+
)
61+
62+
# Finish transaction
63+
op.execute("COMMIT;")

src/codegate/cli.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,11 @@
1414
from codegate.ca.codegate_ca import CertificateAuthority
1515
from codegate.codegate_logging import LogFormat, LogLevel, setup_logging
1616
from codegate.config import Config, ConfigurationError
17-
from codegate.db.connection import init_db_sync, init_session_if_not_exists
17+
from codegate.db.connection import (
18+
init_db_sync,
19+
init_session_if_not_exists,
20+
init_instance,
21+
)
1822
from codegate.pipeline.factory import PipelineFactory
1923
from codegate.pipeline.sensitive_data.manager import SensitiveDataManager
2024
from codegate.providers import crud as provendcrud
@@ -318,6 +322,7 @@ def serve( # noqa: C901
318322
logger = structlog.get_logger("codegate").bind(origin="cli")
319323

320324
init_db_sync(cfg.db_path)
325+
init_instance(cfg.db_path)
321326
init_session_if_not_exists(cfg.db_path)
322327

323328
# Check certificates and create CA if necessary

src/codegate/db/connection.py

Lines changed: 44 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import asyncio
2+
import datetime
23
import json
34
import sqlite3
45
import uuid
@@ -23,6 +24,7 @@
2324
Alert,
2425
GetPromptWithOutputsRow,
2526
GetWorkspaceByNameConditions,
27+
Instance,
2628
IntermediatePromptWithOutputUsageAlerts,
2729
MuxRule,
2830
Output,
@@ -570,6 +572,27 @@ async def add_persona(self, persona: PersonaEmbedding) -> None:
570572
logger.debug(f"Exception type: {type(e)}")
571573
raise AlreadyExistsError(f"Persona '{persona.name}' already exists.")
572574

575+
async def init_instance(self) -> None:
576+
"""
577+
Initializes instance details in the database.
578+
"""
579+
sql = text(
580+
"""
581+
INSERT INTO instance (id, created_at)
582+
VALUES (:id, :created_at)
583+
"""
584+
)
585+
586+
try:
587+
instance = Instance(
588+
id=str(uuid.uuid4()),
589+
created_at=datetime.datetime.now(datetime.timezone.utc),
590+
)
591+
await self._execute_with_no_return(sql, instance.model_dump())
592+
except IntegrityError as e:
593+
logger.debug(f"Exception type: {type(e)}")
594+
raise AlreadyExistsError(f"Instance already initialized.")
595+
573596

574597
class DbReader(DbCodeGate):
575598
def __init__(self, sqlite_path: Optional[str] = None, *args, **kwargs):
@@ -1045,6 +1068,13 @@ async def get_distance_to_persona(
10451068
)
10461069
return persona_distance[0]
10471070

1071+
async def get_instance(self) -> Instance:
1072+
"""
1073+
Get the details of the instance.
1074+
"""
1075+
sql = text("SELECT id, created_at FROM instance")
1076+
return await self._execute_select_pydantic_model(Instance, sql)
1077+
10481078

10491079
class DbTransaction:
10501080
def __init__(self):
@@ -1095,8 +1125,6 @@ def init_db_sync(db_path: Optional[str] = None):
10951125

10961126

10971127
def init_session_if_not_exists(db_path: Optional[str] = None):
1098-
import datetime
1099-
11001128
db_reader = DbReader(db_path)
11011129
sessions = asyncio.run(db_reader.get_sessions())
11021130
# If there are no sessions, create a new one
@@ -1116,5 +1144,19 @@ def init_session_if_not_exists(db_path: Optional[str] = None):
11161144
logger.info("Session in DB initialized successfully.")
11171145

11181146

1147+
def init_instance(db_path: Optional[str] = None):
1148+
db_reader = DbReader(db_path)
1149+
instance = asyncio.run(db_reader.get_instance())
1150+
# Initialize instance if not already initialized.
1151+
if not instance:
1152+
db_recorder = DbRecorder(db_path)
1153+
try:
1154+
asyncio.run(db_recorder.init_instance())
1155+
except Exception as e:
1156+
logger.error(f"Failed to initialize instance in DB: {e}")
1157+
raise
1158+
logger.info("Instance initialized successfully.")
1159+
1160+
11191161
if __name__ == "__main__":
11201162
init_db_sync()

src/codegate/db/models.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,11 @@ class Session(BaseModel):
120120
last_update: datetime.datetime
121121

122122

123+
class Instance(BaseModel):
124+
id: str
125+
created_at: datetime.datetime
126+
127+
123128
# Models for select queries
124129

125130

0 commit comments

Comments
 (0)