Xiaoxia AI
108 lines
4.2 KiB
Python
108 lines
4.2 KiB
Python
"""
|
|
PostgreSQL Workspace Repository 实现
|
|
"""
|
|
|
|
from typing import Optional
|
|
|
|
import psycopg2
|
|
from psycopg2.extras import RealDictCursor
|
|
|
|
from packages.domain.entities import Workspace
|
|
from packages.ports.workspace_repository import WorkspaceRepository
|
|
|
|
|
|
class PostgresWorkspaceRepository(WorkspaceRepository):
|
|
"""Workspace 仓储 PostgreSQL 实现"""
|
|
|
|
def __init__(self, connection_string: str):
|
|
self.connection_string = connection_string
|
|
|
|
def _get_connection(self):
|
|
"""获取数据库连接（使用连接池）"""
|
|
from packages.adapters.postgres.connection_pool import PooledConnection
|
|
|
|
return PooledConnection()
|
|
|
|
def save(self, workspace: Workspace) -> None:
|
|
"""保存工作空间"""
|
|
conn = self._get_connection()
|
|
try:
|
|
with conn.cursor() as cur:
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO workspaces (
|
|
id, name, owner_user_id, subscription_plan,
|
|
subscription_status, subscription_expires_at,
|
|
max_projects, max_storage_gb, used_storage_gb, created_at
|
|
) VALUES (
|
|
%(id)s, %(name)s, %(owner_user_id)s, %(subscription_plan)s,
|
|
%(subscription_status)s, %(subscription_expires_at)s,
|
|
%(max_projects)s, %(max_storage_gb)s, %(used_storage_gb)s, %(created_at)s
|
|
)
|
|
ON CONFLICT (id) DO UPDATE SET
|
|
name = EXCLUDED.name,
|
|
subscription_plan = EXCLUDED.subscription_plan,
|
|
subscription_status = EXCLUDED.subscription_status,
|
|
subscription_expires_at = EXCLUDED.subscription_expires_at,
|
|
max_projects = EXCLUDED.max_projects,
|
|
max_storage_gb = EXCLUDED.max_storage_gb,
|
|
used_storage_gb = EXCLUDED.used_storage_gb
|
|
""",
|
|
{
|
|
"id": workspace.id,
|
|
"name": workspace.name,
|
|
"owner_user_id": workspace.owner_user_id,
|
|
"subscription_plan": workspace.subscription_plan,
|
|
"subscription_status": workspace.subscription_status,
|
|
"subscription_expires_at": workspace.subscription_expires_at,
|
|
"max_projects": workspace.max_projects,
|
|
"max_storage_gb": workspace.max_storage_gb,
|
|
"used_storage_gb": workspace.used_storage_gb,
|
|
"created_at": workspace.created_at,
|
|
},
|
|
)
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
def find_by_id(self, workspace_id: str) -> Optional[Workspace]:
|
|
"""根据 ID 查找工作空间"""
|
|
conn = self._get_connection()
|
|
try:
|
|
with conn.cursor() as cur:
|
|
cur.execute("SELECT * FROM workspaces WHERE id = %s", (workspace_id,))
|
|
row = cur.fetchone()
|
|
|
|
if row:
|
|
return self._row_to_workspace(row)
|
|
return None
|
|
finally:
|
|
conn.close()
|
|
|
|
def delete(self, workspace_id: str) -> bool:
|
|
"""删除工作空间"""
|
|
conn = self._get_connection()
|
|
try:
|
|
with conn.cursor() as cur:
|
|
cur.execute("DELETE FROM workspaces WHERE id = %s", (workspace_id,))
|
|
deleted = cur.rowcount > 0
|
|
conn.commit()
|
|
return deleted
|
|
finally:
|
|
conn.close()
|
|
|
|
def _row_to_workspace(self, row: dict) -> Workspace:
|
|
"""将数据库行转换为 Workspace 对象"""
|
|
return Workspace(
|
|
id=row["id"],
|
|
name=row["name"],
|
|
owner_user_id=row["owner_user_id"],
|
|
subscription_plan=row["subscription_plan"],
|
|
subscription_status=row["subscription_status"],
|
|
subscription_expires_at=row["subscription_expires_at"],
|
|
max_projects=row["max_projects"],
|
|
max_storage_gb=row["max_storage_gb"],
|
|
used_storage_gb=float(row["used_storage_gb"]),
|
|
created_at=row["created_at"],
|
|
)
|