add test module
This commit is contained in:
305
sentence_api/repository.py
Normal file
305
sentence_api/repository.py
Normal file
@@ -0,0 +1,305 @@
|
||||
import json
|
||||
import sqlite3
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from .models import SentenceBoundary, SentenceBoundaryDocument
|
||||
from .store import normalize_video_hash
|
||||
|
||||
|
||||
def utc_now() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
class VideoRepository:
|
||||
def __init__(self, database_path: Path):
|
||||
self.database_path = Path(database_path)
|
||||
self.database_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._initialize()
|
||||
|
||||
def _connect(self) -> sqlite3.Connection:
|
||||
connection = sqlite3.connect(self.database_path, timeout=30)
|
||||
connection.row_factory = sqlite3.Row
|
||||
connection.execute("PRAGMA foreign_keys = ON")
|
||||
connection.execute("PRAGMA journal_mode = WAL")
|
||||
return connection
|
||||
|
||||
def _initialize(self) -> None:
|
||||
with self._connect() as connection:
|
||||
connection.executescript(
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS videos (
|
||||
video_hash TEXT PRIMARY KEY,
|
||||
title TEXT NOT NULL,
|
||||
filename TEXT NOT NULL,
|
||||
stored_filename TEXT NOT NULL,
|
||||
content_type TEXT NOT NULL,
|
||||
size_bytes INTEGER NOT NULL,
|
||||
duration_ms INTEGER,
|
||||
language TEXT,
|
||||
status TEXT NOT NULL,
|
||||
error_message TEXT,
|
||||
algorithm_version TEXT,
|
||||
transcription TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS sentences (
|
||||
video_hash TEXT NOT NULL,
|
||||
sentence_index INTEGER NOT NULL,
|
||||
start_ms INTEGER NOT NULL,
|
||||
end_ms INTEGER NOT NULL,
|
||||
text TEXT,
|
||||
language TEXT,
|
||||
reference_speech_duration_ms INTEGER,
|
||||
PRIMARY KEY (video_hash, sentence_index),
|
||||
FOREIGN KEY (video_hash) REFERENCES videos(video_hash) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS attempts (
|
||||
attempt_id TEXT PRIMARY KEY,
|
||||
video_hash TEXT NOT NULL,
|
||||
sentence_index INTEGER NOT NULL,
|
||||
result_json TEXT NOT NULL,
|
||||
audio_filename TEXT,
|
||||
created_at TEXT NOT NULL,
|
||||
FOREIGN KEY (video_hash) REFERENCES videos(video_hash) ON DELETE CASCADE
|
||||
);
|
||||
"""
|
||||
)
|
||||
|
||||
def upsert_upload(
|
||||
self,
|
||||
*,
|
||||
video_hash: str,
|
||||
title: str,
|
||||
filename: str,
|
||||
stored_filename: str,
|
||||
content_type: str,
|
||||
size_bytes: int,
|
||||
language: Optional[str],
|
||||
) -> None:
|
||||
normalized_hash = normalize_video_hash(video_hash)
|
||||
now = utc_now()
|
||||
with self._connect() as connection:
|
||||
existing = connection.execute(
|
||||
"SELECT created_at FROM videos WHERE video_hash = ?",
|
||||
(normalized_hash,),
|
||||
).fetchone()
|
||||
created_at = existing["created_at"] if existing else now
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO videos (
|
||||
video_hash, title, filename, stored_filename, content_type,
|
||||
size_bytes, duration_ms, language, status, error_message,
|
||||
algorithm_version, transcription, created_at, updated_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?, NULL, ?, 'uploaded', NULL, NULL, NULL, ?, ?)
|
||||
ON CONFLICT(video_hash) DO UPDATE SET
|
||||
title = excluded.title,
|
||||
filename = excluded.filename,
|
||||
stored_filename = excluded.stored_filename,
|
||||
content_type = excluded.content_type,
|
||||
size_bytes = excluded.size_bytes,
|
||||
language = excluded.language,
|
||||
status = 'uploaded',
|
||||
error_message = NULL,
|
||||
updated_at = excluded.updated_at
|
||||
""",
|
||||
(
|
||||
normalized_hash,
|
||||
title,
|
||||
filename,
|
||||
stored_filename,
|
||||
content_type,
|
||||
size_bytes,
|
||||
language,
|
||||
created_at,
|
||||
now,
|
||||
),
|
||||
)
|
||||
|
||||
def mark_processing(self, video_hash: str) -> None:
|
||||
self._update_status(video_hash, "processing", None)
|
||||
|
||||
def mark_failed(self, video_hash: str, message: str) -> None:
|
||||
self._update_status(video_hash, "failed", message[:4000])
|
||||
|
||||
def _update_status(self, video_hash: str, status: str, error_message: Optional[str]) -> None:
|
||||
with self._connect() as connection:
|
||||
connection.execute(
|
||||
"UPDATE videos SET status = ?, error_message = ?, updated_at = ? WHERE video_hash = ?",
|
||||
(status, error_message, utc_now(), normalize_video_hash(video_hash)),
|
||||
)
|
||||
|
||||
def save_processing_result(
|
||||
self,
|
||||
document: SentenceBoundaryDocument,
|
||||
transcription: Optional[str],
|
||||
) -> None:
|
||||
with self._connect() as connection:
|
||||
connection.execute(
|
||||
"DELETE FROM sentences WHERE video_hash = ?",
|
||||
(document.video_hash,),
|
||||
)
|
||||
connection.executemany(
|
||||
"""
|
||||
INSERT INTO sentences (
|
||||
video_hash, sentence_index, start_ms, end_ms, text,
|
||||
language, reference_speech_duration_ms
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
[
|
||||
(
|
||||
document.video_hash,
|
||||
sentence.index,
|
||||
sentence.start_ms,
|
||||
sentence.end_ms,
|
||||
sentence.text,
|
||||
sentence.language,
|
||||
sentence.reference_speech_duration_ms,
|
||||
)
|
||||
for sentence in document.sentences
|
||||
],
|
||||
)
|
||||
connection.execute(
|
||||
"""
|
||||
UPDATE videos SET duration_ms = ?, status = 'ready', error_message = NULL,
|
||||
algorithm_version = ?, transcription = ?, updated_at = ?
|
||||
WHERE video_hash = ?
|
||||
""",
|
||||
(
|
||||
document.duration_ms,
|
||||
document.algorithm_version,
|
||||
transcription,
|
||||
utc_now(),
|
||||
document.video_hash,
|
||||
),
|
||||
)
|
||||
|
||||
def list_videos(self) -> List[Dict[str, Any]]:
|
||||
with self._connect() as connection:
|
||||
rows = connection.execute(
|
||||
"""
|
||||
SELECT v.*, COUNT(s.sentence_index) AS sentence_count
|
||||
FROM videos v
|
||||
LEFT JOIN sentences s ON s.video_hash = v.video_hash
|
||||
GROUP BY v.video_hash
|
||||
ORDER BY v.created_at DESC
|
||||
"""
|
||||
).fetchall()
|
||||
return [dict(row) for row in rows]
|
||||
|
||||
def get_video(self, video_hash: str) -> Optional[Dict[str, Any]]:
|
||||
with self._connect() as connection:
|
||||
row = connection.execute(
|
||||
"""
|
||||
SELECT v.*, COUNT(s.sentence_index) AS sentence_count
|
||||
FROM videos v
|
||||
LEFT JOIN sentences s ON s.video_hash = v.video_hash
|
||||
WHERE v.video_hash = ?
|
||||
GROUP BY v.video_hash
|
||||
""",
|
||||
(normalize_video_hash(video_hash),),
|
||||
).fetchone()
|
||||
return dict(row) if row else None
|
||||
|
||||
def get_document(self, video_hash: str) -> Optional[SentenceBoundaryDocument]:
|
||||
video = self.get_video(video_hash)
|
||||
if video is None or video["duration_ms"] is None:
|
||||
return None
|
||||
with self._connect() as connection:
|
||||
rows = connection.execute(
|
||||
"SELECT * FROM sentences WHERE video_hash = ? ORDER BY sentence_index",
|
||||
(normalize_video_hash(video_hash),),
|
||||
).fetchall()
|
||||
return SentenceBoundaryDocument(
|
||||
video_hash=video["video_hash"],
|
||||
duration_ms=video["duration_ms"],
|
||||
algorithm_version=video["algorithm_version"] or "unknown",
|
||||
sentences=[
|
||||
SentenceBoundary(
|
||||
index=row["sentence_index"],
|
||||
start_ms=row["start_ms"],
|
||||
end_ms=row["end_ms"],
|
||||
text=row["text"],
|
||||
language=row["language"],
|
||||
reference_speech_duration_ms=row["reference_speech_duration_ms"],
|
||||
)
|
||||
for row in rows
|
||||
],
|
||||
)
|
||||
|
||||
def update_sentence_text(
|
||||
self,
|
||||
video_hash: str,
|
||||
sentence_index: int,
|
||||
text: str,
|
||||
language: Optional[str],
|
||||
) -> Optional[SentenceBoundary]:
|
||||
normalized_hash = normalize_video_hash(video_hash)
|
||||
with self._connect() as connection:
|
||||
cursor = connection.execute(
|
||||
"""
|
||||
UPDATE sentences SET text = ?, language = COALESCE(?, language)
|
||||
WHERE video_hash = ? AND sentence_index = ?
|
||||
""",
|
||||
(text.strip(), language, normalized_hash, sentence_index),
|
||||
)
|
||||
if cursor.rowcount == 0:
|
||||
return None
|
||||
connection.execute(
|
||||
"UPDATE videos SET updated_at = ? WHERE video_hash = ?",
|
||||
(utc_now(), normalized_hash),
|
||||
)
|
||||
row = connection.execute(
|
||||
"SELECT * FROM sentences WHERE video_hash = ? AND sentence_index = ?",
|
||||
(normalized_hash, sentence_index),
|
||||
).fetchone()
|
||||
return SentenceBoundary(
|
||||
index=row["sentence_index"],
|
||||
start_ms=row["start_ms"],
|
||||
end_ms=row["end_ms"],
|
||||
text=row["text"],
|
||||
language=row["language"],
|
||||
reference_speech_duration_ms=row["reference_speech_duration_ms"],
|
||||
)
|
||||
|
||||
def record_attempt(
|
||||
self,
|
||||
*,
|
||||
attempt_id: str,
|
||||
video_hash: str,
|
||||
sentence_index: int,
|
||||
result: Dict[str, Any],
|
||||
audio_filename: Optional[str],
|
||||
) -> None:
|
||||
with self._connect() as connection:
|
||||
connection.execute(
|
||||
"""
|
||||
INSERT INTO attempts (
|
||||
attempt_id, video_hash, sentence_index, result_json,
|
||||
audio_filename, created_at
|
||||
) VALUES (?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
attempt_id,
|
||||
normalize_video_hash(video_hash),
|
||||
sentence_index,
|
||||
json.dumps(result, ensure_ascii=False),
|
||||
audio_filename,
|
||||
utc_now(),
|
||||
),
|
||||
)
|
||||
|
||||
def delete_video(self, video_hash: str) -> Optional[Dict[str, Any]]:
|
||||
video = self.get_video(video_hash)
|
||||
if video is None:
|
||||
return None
|
||||
with self._connect() as connection:
|
||||
connection.execute(
|
||||
"DELETE FROM videos WHERE video_hash = ?",
|
||||
(normalize_video_hash(video_hash),),
|
||||
)
|
||||
return video
|
||||
Reference in New Issue
Block a user