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

Commit 70b2d6e

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 70b2d6e

File tree

4 files changed

+117
-3
lines changed

4 files changed

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