Files
tts/server/db.py
2026-02-25 19:04:18 +09:00

160 lines
4.2 KiB
Python

import os
from typing import List, Optional, Dict, Any
import psycopg2
import psycopg2.extras
def get_conn():
user = os.getenv("DB_USER")
password = os.getenv("DB_PASSWORD")
if not user or not password:
raise RuntimeError("DB_USER 또는 DB_PASSWORD가 설정되지 않았습니다.")
return psycopg2.connect(
host="ncue.net",
port=5432,
dbname="tts",
user=user,
password=password,
)
def init_db():
with get_conn() as conn:
with conn.cursor() as cur:
cur.execute(
"""
CREATE TABLE IF NOT EXISTS tts_items (
id SERIAL PRIMARY KEY,
text TEXT NOT NULL,
filename TEXT,
size_bytes BIGINT,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
"""
)
cur.execute(
"""
ALTER TABLE tts_items
ADD COLUMN IF NOT EXISTS size_bytes BIGINT;
"""
)
cur.execute(
"""
CREATE INDEX IF NOT EXISTS tts_items_created_at_idx
ON tts_items (created_at DESC);
"""
)
conn.commit()
def create_item(text: str) -> Dict[str, Any]:
with get_conn() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
cur.execute(
"""
INSERT INTO tts_items (text)
VALUES (%s)
RETURNING id, created_at;
""",
(text,),
)
row = cur.fetchone()
conn.commit()
return row
def update_filename(tts_id: int, filename: str) -> None:
with get_conn() as conn:
with conn.cursor() as cur:
cur.execute(
"""
UPDATE tts_items
SET filename = %s
WHERE id = %s;
""",
(filename, tts_id),
)
conn.commit()
def update_size_bytes(tts_id: int, size_bytes: int) -> None:
with get_conn() as conn:
with conn.cursor() as cur:
cur.execute(
"""
UPDATE tts_items
SET size_bytes = %s
WHERE id = %s;
""",
(size_bytes, tts_id),
)
conn.commit()
def list_items() -> List[Dict[str, Any]]:
with get_conn() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
cur.execute(
"""
SELECT id, created_at, filename, size_bytes
FROM tts_items
ORDER BY created_at DESC;
"""
)
rows = cur.fetchall()
return rows
def get_item(tts_id: int) -> Optional[Dict[str, Any]]:
with get_conn() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
cur.execute(
"""
SELECT id, text, filename, size_bytes, created_at
FROM tts_items
WHERE id = %s;
""",
(tts_id,),
)
row = cur.fetchone()
return row
def delete_items(ids: List[int]) -> List[Dict[str, Any]]:
with get_conn() as conn:
with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
cur.execute(
"""
SELECT id, filename
FROM tts_items
WHERE id = ANY(%s);
""",
(ids,),
)
rows = cur.fetchall()
cur.execute(
"""
DELETE FROM tts_items
WHERE id = ANY(%s);
""",
(ids,),
)
conn.commit()
return rows
def delete_item_by_id(tts_id: int) -> None:
with get_conn() as conn:
with conn.cursor() as cur:
cur.execute(
"""
DELETE FROM tts_items
WHERE id = %s;
""",
(tts_id,),
)
conn.commit()