160 lines
4.2 KiB
Python
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()
|