Files
zhuangxiu/services/api/app/repositories/postgres.py
T

85 lines
3.0 KiB
Python

from __future__ import annotations
import json
from functools import lru_cache
import psycopg
from psycopg.types.json import Jsonb
from app.domain.models import ProjectSnapshot
from app.repositories.memory import create_demo_project
class PostgresProjectRepository:
def __init__(self, database_url: str) -> None:
self.database_url = database_url
self.ensure_schema()
if self.get("demo-apartment") is None:
self.save(create_demo_project())
def _connect(self):
return psycopg.connect(self.database_url)
def ensure_schema(self) -> None:
with self._connect() as connection, connection.cursor() as cursor:
cursor.execute(
"""
CREATE TABLE IF NOT EXISTS projects (
project_id TEXT PRIMARY KEY,
name TEXT NOT NULL,
stage TEXT NOT NULL,
revision INTEGER NOT NULL,
snapshot JSONB NOT NULL,
updated_at TIMESTAMPTZ NOT NULL
)
"""
)
def get(self, project_id: str) -> ProjectSnapshot | None:
with self._connect() as connection, connection.cursor() as cursor:
cursor.execute("SELECT snapshot FROM projects WHERE project_id = %s", (project_id,))
row = cursor.fetchone()
if row is None:
return None
payload = json.loads(row[0]) if isinstance(row[0], str) else row[0]
return ProjectSnapshot.model_validate(payload)
def list(self) -> list[ProjectSnapshot]:
with self._connect() as connection, connection.cursor() as cursor:
cursor.execute("SELECT snapshot FROM projects ORDER BY updated_at DESC")
rows = cursor.fetchall()
return [
ProjectSnapshot.model_validate(json.loads(row[0]) if isinstance(row[0], str) else row[0])
for row in rows
]
def save(self, project: ProjectSnapshot) -> ProjectSnapshot:
payload = project.model_dump(mode="json")
with self._connect() as connection, connection.cursor() as cursor:
cursor.execute(
"""
INSERT INTO projects (project_id, name, stage, revision, snapshot, updated_at)
VALUES (%s, %s, %s, %s, %s, %s)
ON CONFLICT (project_id) DO UPDATE SET
name = EXCLUDED.name,
stage = EXCLUDED.stage,
revision = EXCLUDED.revision,
snapshot = EXCLUDED.snapshot,
updated_at = EXCLUDED.updated_at
""",
(
project.project_id,
project.name,
project.stage.value,
project.revision,
Jsonb(payload),
project.updated_at,
),
)
return project
@lru_cache(maxsize=4)
def postgres_repository(database_url: str) -> PostgresProjectRepository:
return PostgresProjectRepository(database_url)