Files
copykar/db/database.py
T

86 lines
2.7 KiB
Python

import asyncpg
import os
from typing import Optional
DATABASE_URL = os.getenv(
"DATABASE_URL",
f"postgresql://{os.getenv('POSTGRES_USER', 'postgres')}:{os.getenv('POSTGRES_PASSWORD', 'postgres')}@{os.getenv('POSTGRES_HOST', 'localhost')}:{os.getenv('POSTGRES_PORT', '5432')}/{os.getenv('POSTGRES_DB', 'copykar')}"
)
SCHEMA = """
CREATE TABLE IF NOT EXISTS sources (
id SERIAL PRIMARY KEY,
channel_id BIGINT UNIQUE NOT NULL,
username VARCHAR(255),
title VARCHAR(255),
is_active BOOLEAN DEFAULT TRUE,
created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS targets (
id SERIAL PRIMARY KEY,
channel_id BIGINT UNIQUE NOT NULL,
title VARCHAR(255),
username VARCHAR(255),
post_interval_min INT DEFAULT 30,
last_post_time TIMESTAMPTZ,
is_active BOOLEAN DEFAULT TRUE,
created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS posts (
id BIGSERIAL PRIMARY KEY,
source_channel_id BIGINT NOT NULL,
source_message_id BIGINT NOT NULL,
raw_text TEXT,
media_path TEXT,
media_type VARCHAR(64),
content_hash VARCHAR(128),
tags TEXT[] DEFAULT '{}',
is_duplicate BOOLEAN DEFAULT FALSE,
duplicate_of_id BIGINT REFERENCES posts(id) ON DELETE SET NULL,
similarity_reason TEXT,
subject VARCHAR(255),
ai_text TEXT,
suggested_target_id INT REFERENCES targets(id) ON DELETE SET NULL,
target_channel_id INT REFERENCES targets(id) ON DELETE SET NULL,
status VARCHAR(32) DEFAULT 'pending_ai',
review_message_id BIGINT,
scheduled_at TIMESTAMPTZ,
published_at TIMESTAMPTZ,
created_at TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP,
CONSTRAINT unique_source_message UNIQUE (source_channel_id, source_message_id)
);
CREATE INDEX IF NOT EXISTS idx_posts_content_hash ON posts(content_hash);
CREATE INDEX IF NOT EXISTS idx_posts_status ON posts(status);
CREATE INDEX IF NOT EXISTS idx_posts_target_status ON posts(target_channel_id, status);
CREATE INDEX IF NOT EXISTS idx_posts_tags ON posts USING GIN (tags);
CREATE INDEX IF NOT EXISTS idx_posts_created_at ON posts(created_at DESC);
CREATE TABLE IF NOT EXISTS settings (
key VARCHAR(128) PRIMARY KEY,
value TEXT NOT NULL,
description TEXT
);
"""
_pool: Optional[asyncpg.Pool] = None
async def get_db_pool(dsn: str = DATABASE_URL) -> asyncpg.Pool:
global _pool
if _pool is None or _pool._closed:
_pool = await asyncpg.create_pool(dsn=dsn, min_size=2, max_size=10)
return _pool
async def close_db_pool():
global _pool
if _pool is not None and not _pool._closed:
await _pool.close()
_pool = None
async def init_db(dsn: str = DATABASE_URL):
pool = await get_db_pool(dsn)
async with pool.acquire() as conn:
await conn.execute(SCHEMA)