src/database/seed.py - master - asigatchov/vb-ai-api

src/database/seed.py
161 строка4 KB
Alexander Sigatchov
init
09 фев 2026, 08:49
09 фев 2026, 08:49b989369
from __future__ import annotations

import json
import re
from pathlib import Path

from sqlalchemy import select
from sqlalchemy.orm import Session

from src.config.security import hash_password
from src.database.models import Annotation, Folder, Project, Rally, User


def _extract_export_value(source: str, const_name: str) -> str:
    marker = f"export const {const_name}"
    idx = source.find(marker)
    if idx == -1:
        raise ValueError(f"Constant {const_name} not found in db-data.js")

    eq_idx = source.find("=", idx)
    if eq_idx == -1:
        raise ValueError(f"Malformed export for {const_name}")

    i = eq_idx + 1
    while i < len(source) and source[i].isspace():
        i += 1

    out: list[str] = []
    depth = 0
    quote: str | None = None
    escape = False

    while i < len(source):
        ch = source[i]
        if quote:
            out.append(ch)
            if escape:
                escape = False
            elif ch == "\\":
                escape = True
            elif ch == quote:
                quote = None
        else:
            if ch in ('"', "'"):
                quote = ch
                out.append(ch)
            elif ch in "[{":
                depth += 1
                out.append(ch)
            elif ch in "]}":
                depth -= 1
                out.append(ch)
            elif ch == ";" and depth == 0:
                break
            else:
                out.append(ch)
        i += 1

    return "".join(out).strip()


def _to_json_literal(js_literal: str) -> str:
    text = re.sub(r"//.*", "", js_literal)
    text = text.replace("'", '"')
    text = re.sub(r"([\{,]\s*)([A-Za-z_][A-Za-z0-9_]*)(\s*:)", r'\1"\2"\3', text)
    text = re.sub(r"\btrue\b", "true", text)
    text = re.sub(r"\bfalse\b", "false", text)
    text = re.sub(r"\bnull\b", "null", text)
    text = re.sub(r",\s*([}\]])", r"\1", text)
    return text


def _load_mock_data(db_data_path: Path) -> dict[str, list | dict]:
    raw = db_data_path.read_text(encoding="utf-8")

    keys = [
        "INITIAL_KEYPOINTS",
        "MOCK_RALLIES",
        "MOCK_PROJECTS",
        "MOCK_FOLDERS",
        "MOCK_USERS",
    ]
    parsed: dict[str, list | dict] = {}
    for key in keys:
        value = _extract_export_value(raw, key)
        parsed[key] = json.loads(_to_json_literal(value))
    return parsed


def seed_initial_data(db: Session, db_data_path: Path) -> None:
    has_users = db.scalar(select(User.id).limit(1))
    if has_users:
        return

    data = _load_mock_data(db_data_path)

    users = [
        User(
            id=item["id"],
            name=item["name"],
            email=item["email"],
            avatar=item.get("avatar"),
            password_hash=hash_password("password123"),
        )
        for item in data["MOCK_USERS"]
    ]
    db.add_all(users)

    owner_id = users[0].id

    folders: list[Folder] = []
    folder_map: dict[str, str] = {}
    for idx, folder_name in enumerate(data["MOCK_FOLDERS"], start=1):
        folder_id = f"folder_{idx}"
        folder_map[folder_name] = folder_id
        folders.append(Folder(id=folder_id, user_id=owner_id, name=folder_name, created_at="2024-05-01T00:00:00+00:00"))
    db.add_all(folders)

    projects: list[Project] = []
    for item in data["MOCK_PROJECTS"]:
        projects.append(
            Project(
                id=item["id"],
                user_id=owner_id,
                folder_id=folder_map.get(item.get("folder", "")),
                title=item["title"],
                date=item["date"],
                status=item["status"],
                video_src=item["videoSrc"],
                thumbnail_url=None,
            )
        )
    db.add_all(projects)

    keypoints = data["INITIAL_KEYPOINTS"]
    annotated_projects = [project for project in projects if project.status == "annotated"]
    for project in annotated_projects:
        db.add(
            Annotation(
                project_id=project.id,
                keypoints_json=json.dumps(keypoints),
                bbox_json=json.dumps([100, 100, 1024, 576]),
                image_width=1280,
                image_height=720,
            )
        )

    rallies = [
        Rally(
            id=item["id"],
            project_id=projects[0].id,
            start_time=float(item["startTime"]),
            end_time=float(item["endTime"]),
            type=item["type"],
            description=item["description"],
        )
        for item in data["MOCK_RALLIES"]
    ]
    db.add_all(rallies)

    db.commit()