345 lines
9.4 KiB
Python
345 lines
9.4 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
from datetime import UTC, date, datetime
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import psycopg
|
|
from psycopg.rows import dict_row
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(PROJECT_ROOT))
|
|
|
|
from db import MIGRATIONS, SCHEMA_STATEMENTS, _normalize_dsn
|
|
|
|
REQUEST_COLUMNS = (
|
|
"id",
|
|
"user_id",
|
|
"username",
|
|
"timestamp",
|
|
"prompt_text",
|
|
"enhanced_prompt",
|
|
"input_file_path",
|
|
"output_file_path",
|
|
"status",
|
|
"comfy_prompt_id",
|
|
"payment_charge_id",
|
|
"payment_refunded",
|
|
"is_premium",
|
|
"mode",
|
|
"seed",
|
|
"error_text",
|
|
"finished_at",
|
|
)
|
|
|
|
USER_COLUMNS = (
|
|
"user_id",
|
|
"username",
|
|
"credits",
|
|
"last_daily_credit",
|
|
"referrer_id",
|
|
"first_paid_at",
|
|
"created_at",
|
|
"is_admin",
|
|
"is_blocked",
|
|
"age_confirmed",
|
|
"age_declined",
|
|
)
|
|
|
|
BLOCKED_COLUMNS = (
|
|
"id",
|
|
"user_id",
|
|
"timestamp",
|
|
"prompt",
|
|
"reason",
|
|
"mode",
|
|
"has_person",
|
|
)
|
|
|
|
PAYMENT_COLUMNS = (
|
|
"payload",
|
|
"user_id",
|
|
"timestamp",
|
|
"mode",
|
|
"prompt",
|
|
"username",
|
|
"input_file_path",
|
|
"amount",
|
|
"status",
|
|
"is_premium",
|
|
"telegram_payment_charge_id",
|
|
"request_id",
|
|
)
|
|
|
|
SETTING_COLUMNS = ("key", "value")
|
|
|
|
|
|
def parse_dt(value: Any) -> datetime | None:
|
|
if value in (None, ""):
|
|
return None
|
|
if isinstance(value, datetime):
|
|
parsed = value
|
|
else:
|
|
parsed = datetime.fromisoformat(str(value))
|
|
if parsed.tzinfo is None:
|
|
parsed = parsed.replace(tzinfo=UTC)
|
|
return parsed
|
|
|
|
|
|
def parse_day(value: Any) -> date | None:
|
|
if value in (None, ""):
|
|
return None
|
|
if isinstance(value, datetime):
|
|
return value.date()
|
|
if isinstance(value, date):
|
|
return value
|
|
return date.fromisoformat(str(value))
|
|
|
|
|
|
def parse_bool(value: Any) -> bool:
|
|
return bool(int(value or 0))
|
|
|
|
|
|
def table_names(connection: sqlite3.Connection) -> set[str]:
|
|
rows = connection.execute(
|
|
"SELECT name FROM sqlite_master WHERE type = 'table'"
|
|
).fetchall()
|
|
return {str(row["name"]) for row in rows}
|
|
|
|
|
|
def read_rows(
|
|
connection: sqlite3.Connection,
|
|
table: str,
|
|
columns: tuple[str, ...],
|
|
defaults: dict[str, Any],
|
|
) -> list[dict[str, Any]]:
|
|
if table not in table_names(connection):
|
|
print(f" {table}: в SQLite нет таблицы, пропуск")
|
|
return []
|
|
existing = {
|
|
str(row["name"])
|
|
for row in connection.execute(f"PRAGMA table_info({table})")
|
|
}
|
|
select = ", ".join(
|
|
column if column in existing else f":{column} AS {column}"
|
|
for column in columns
|
|
)
|
|
params = {
|
|
column: defaults[column]
|
|
for column in columns
|
|
if column not in existing
|
|
}
|
|
query = f"SELECT {select} FROM {table}"
|
|
cursor = (
|
|
connection.execute(query, params)
|
|
if params
|
|
else connection.execute(query)
|
|
)
|
|
return [dict(row) for row in cursor.fetchall()]
|
|
|
|
|
|
def insert_rows(
|
|
connection: psycopg.Connection,
|
|
table: str,
|
|
columns: tuple[str, ...],
|
|
rows: list[dict[str, Any]],
|
|
conflict: str | None,
|
|
) -> int:
|
|
if not rows:
|
|
return 0
|
|
placeholders = ", ".join(["%s"] * len(columns))
|
|
conflict_clause = (
|
|
f"ON CONFLICT ({conflict}) DO NOTHING"
|
|
if conflict
|
|
else "ON CONFLICT DO NOTHING"
|
|
)
|
|
query = (
|
|
f"INSERT INTO {table} ({', '.join(columns)}) "
|
|
f"VALUES ({placeholders}) {conflict_clause}"
|
|
)
|
|
with connection.cursor() as cursor:
|
|
cursor.executemany(
|
|
query,
|
|
[tuple(row[column] for column in columns) for row in rows],
|
|
)
|
|
return cursor.rowcount
|
|
|
|
|
|
def ensure_schema(connection: psycopg.Connection) -> None:
|
|
with connection.cursor() as cursor:
|
|
for statement in (*SCHEMA_STATEMENTS, *MIGRATIONS):
|
|
cursor.execute(statement)
|
|
|
|
|
|
def reset_sequence(connection: psycopg.Connection, table: str) -> None:
|
|
with connection.cursor() as cursor:
|
|
cursor.execute(
|
|
f"""
|
|
SELECT setval(
|
|
pg_get_serial_sequence('{table}', 'id'),
|
|
COALESCE((SELECT MAX(id) FROM {table}), 1),
|
|
(SELECT MAX(id) IS NOT NULL FROM {table})
|
|
)
|
|
"""
|
|
)
|
|
|
|
|
|
def migrate(sqlite_path: Path, database_url: str) -> None:
|
|
if not sqlite_path.exists():
|
|
raise SystemExit(f"SQLite-база не найдена: {sqlite_path}")
|
|
|
|
source = sqlite3.connect(
|
|
f"{sqlite_path.resolve().as_uri()}?mode=ro", uri=True
|
|
)
|
|
source.row_factory = sqlite3.Row
|
|
try:
|
|
users = read_rows(
|
|
source,
|
|
"users",
|
|
USER_COLUMNS,
|
|
{
|
|
"username": None,
|
|
"credits": 3,
|
|
"last_daily_credit": date.today().isoformat(),
|
|
"referrer_id": None,
|
|
"first_paid_at": None,
|
|
"created_at": datetime.now().astimezone().isoformat(),
|
|
"is_admin": 0,
|
|
"is_blocked": 0,
|
|
"age_confirmed": 0,
|
|
"age_declined": 0,
|
|
},
|
|
)
|
|
requests = read_rows(
|
|
source,
|
|
"requests",
|
|
REQUEST_COLUMNS,
|
|
{
|
|
"username": None,
|
|
"enhanced_prompt": None,
|
|
"payment_charge_id": None,
|
|
"payment_refunded": 0,
|
|
"is_premium": 0,
|
|
"seed": None,
|
|
"error_text": None,
|
|
"finished_at": None,
|
|
},
|
|
)
|
|
blocked = read_rows(
|
|
source,
|
|
"blocked_requests",
|
|
BLOCKED_COLUMNS,
|
|
{"mode": None, "has_person": 0},
|
|
)
|
|
payments = read_rows(
|
|
source,
|
|
"pending_payments",
|
|
PAYMENT_COLUMNS,
|
|
{"username": None, "input_file_path": None, "is_premium": 0},
|
|
)
|
|
settings = read_rows(source, "app_settings", SETTING_COLUMNS, {})
|
|
finally:
|
|
source.close()
|
|
|
|
for row in users:
|
|
row["last_daily_credit"] = parse_day(row["last_daily_credit"])
|
|
row["first_paid_at"] = parse_dt(row["first_paid_at"])
|
|
row["created_at"] = parse_dt(row["created_at"])
|
|
row["is_admin"] = parse_bool(row["is_admin"])
|
|
row["is_blocked"] = parse_bool(row["is_blocked"])
|
|
row["age_confirmed"] = parse_bool(row["age_confirmed"])
|
|
row["age_declined"] = parse_bool(row["age_declined"])
|
|
|
|
for row in requests:
|
|
row["timestamp"] = parse_dt(row["timestamp"])
|
|
row["finished_at"] = parse_dt(row["finished_at"])
|
|
row["payment_refunded"] = parse_bool(row["payment_refunded"])
|
|
row["is_premium"] = parse_bool(row["is_premium"])
|
|
|
|
for row in blocked:
|
|
row["timestamp"] = parse_dt(row["timestamp"])
|
|
row["has_person"] = parse_bool(row["has_person"])
|
|
|
|
for row in payments:
|
|
row["timestamp"] = parse_dt(row["timestamp"])
|
|
row["is_premium"] = parse_bool(row["is_premium"])
|
|
|
|
dsn = _normalize_dsn(database_url)
|
|
with psycopg.connect(
|
|
dsn, row_factory=dict_row, autocommit=True
|
|
) as target:
|
|
ensure_schema(target)
|
|
with target.transaction():
|
|
inserted_users = insert_rows(
|
|
target, "users", USER_COLUMNS, users, "user_id"
|
|
)
|
|
inserted_requests = insert_rows(
|
|
target, "requests", REQUEST_COLUMNS, requests, None
|
|
)
|
|
inserted_blocked = insert_rows(
|
|
target,
|
|
"blocked_requests",
|
|
BLOCKED_COLUMNS,
|
|
blocked,
|
|
None,
|
|
)
|
|
inserted_payments = insert_rows(
|
|
target,
|
|
"pending_payments",
|
|
PAYMENT_COLUMNS,
|
|
payments,
|
|
"payload",
|
|
)
|
|
inserted_settings = insert_rows(
|
|
target,
|
|
"app_settings",
|
|
SETTING_COLUMNS,
|
|
settings,
|
|
"key",
|
|
)
|
|
reset_sequence(target, "requests")
|
|
reset_sequence(target, "blocked_requests")
|
|
|
|
print(f"SQLite: {sqlite_path}")
|
|
print("PostgreSQL: подключение выполнено")
|
|
print(f"users: {inserted_users} из {len(users)}")
|
|
print(f"requests: {inserted_requests} из {len(requests)}")
|
|
print(f"blocked_requests: {inserted_blocked} из {len(blocked)}")
|
|
print(f"pending_payments: {inserted_payments} из {len(payments)}")
|
|
print(f"app_settings: {inserted_settings} из {len(settings)}")
|
|
print("Готово.")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(
|
|
description="Перенос данных из SQLite в PostgreSQL"
|
|
)
|
|
parser.add_argument(
|
|
"--sqlite",
|
|
default=os.getenv("SQLITE_PATH", str(PROJECT_ROOT / "data" / "bot.db")),
|
|
help="путь к старой SQLite-базе",
|
|
)
|
|
parser.add_argument(
|
|
"--database-url",
|
|
default=os.getenv("DATABASE_URL", ""),
|
|
help="строка подключения PostgreSQL",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
if not args.database_url:
|
|
print(
|
|
"DATABASE_URL не задан. Передай --database-url или переменную окружения.",
|
|
file=sys.stderr,
|
|
)
|
|
raise SystemExit(1)
|
|
|
|
migrate(Path(args.sqlite), args.database_url)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|