""" storage/pg_storage.py — Supabase PostgreSQL: quản lý 4 bảng theo ERD Tables: persons — Thông tin cá nhân visits — Lượt ghé thăm face_embeddings — Vector embedding 512 chiều recognition_logs — Nhật ký nhận diện (đầy đủ field theo ERD) """ import json import logging import uuid import psycopg2 from psycopg2.extras import RealDictCursor import config logger = logging.getLogger("pg_storage") class PostgresStorage: """Quản lý tất cả thao tác SQL trên Supabase PostgreSQL.""" def __init__(self, uri: str = config.SUPABASE_DB_URI): self.uri = uri def _conn(self): return psycopg2.connect(self.uri) # ── Load embeddings ────────────────────────────────────────────────── def load_all_embeddings(self) -> list[dict]: """ Nạp toàn bộ active embeddings và metadata người dùng từ DB. Trả về list[{person_id, full_name, embedding, face_crop_mongo_id, embedding_id}] """ query = """ SELECT p.id::text AS person_id, p.full_name, e.id::text AS embedding_id, e.embedding, e.face_crop_mongo_id FROM persons p JOIN face_embeddings e ON p.id = e.person_id WHERE e.is_active = TRUE ORDER BY p.full_name; """ conn = self._conn() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute(query) rows = cur.fetchall() result = [] for r in rows: emb = r["embedding"] if isinstance(emb, str): cleaned = emb.strip("[]") embedding = [float(x) for x in cleaned.split(",") if x.strip()] elif isinstance(emb, list): embedding = [float(x) for x in emb] else: continue result.append({ "person_id": r["person_id"], "full_name": r["full_name"], "embedding_id": r["embedding_id"], "embedding": embedding, "face_crop_mongo_id": r["face_crop_mongo_id"], }) logger.info(f"[PG] Loaded {len(result)} embeddings from Supabase.") return result finally: conn.close() # ── match_face RPC ─────────────────────────────────────────────────── def match_face_rpc( self, embedding: list[float], threshold: float = config.SIMILARITY_THRESHOLD, ) -> dict | None: """ Gọi Supabase RPC function `match_face` để tìm khuôn mặt khớp nhất. Returns: dict với {person_id, full_name, embedding_id, similarity} hoặc None """ emb_str = "[" + ",".join(map(str, embedding)) + "]" conn = self._conn() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute( "SELECT * FROM match_face(%s::vector, %s);", (emb_str, threshold), ) row = cur.fetchone() if row: return dict(row) return None except Exception as e: logger.warning(f"[PG] match_face RPC failed (will fallback to RAM): {e}") return None finally: conn.close() # ── register_new_face ──────────────────────────────────────────────── def register_new_face( self, name: str, embedding: list[float], face_crop_mongo_id: str, model_name: str = config.MODEL_NAME, ) -> tuple[str, str]: """ Đăng ký người mới: thêm vào persons + face_embeddings. Returns: (person_id, embedding_id) """ person_id = str(uuid.uuid4()) emb_id = str(uuid.uuid4()) emb_str = "[" + ",".join(map(str, embedding)) + "]" conn = self._conn() try: with conn.cursor() as cur: cur.execute( """ INSERT INTO persons (id, full_name, status, total_appearances, total_visits, created_at, updated_at) VALUES (%s, %s, 'active', 1, 1, NOW(), NOW()); """, (person_id, name), ) cur.execute( """ INSERT INTO face_embeddings (id, person_id, embedding, model_name, model_version, face_crop_mongo_id, is_primary, is_active, created_at) VALUES (%s, %s, %s, %s, '1.0.0', %s, TRUE, TRUE, NOW()); """, (emb_id, person_id, emb_str, model_name, face_crop_mongo_id), ) conn.commit() logger.info(f"[PG] Registered new face: {name} (person_id={person_id})") return person_id, emb_id except Exception as e: conn.rollback() logger.error(f"[PG] register_new_face failed: {e}") raise finally: conn.close() # ── record_matched_visit ───────────────────────────────────────────── def record_matched_visit( self, person_id: str, embedding_id: str, face_crop_mongo_id: str, raw_image_mongo_id: str, similarity: float, bbox: dict, ui_color: str = "#00DC50", ui_track_index: int = 0, ) -> tuple[str, str]: """ Ghi nhận lượt ghé thăm cho khuôn mặt khớp. - Tạo hoặc cập nhật visit (trong 30 phút) - Ghi recognition_logs đầy đủ field theo ERD - Cập nhật persons (total_appearances, total_visits, last_seen_at) Returns: (visit_id, log_id) """ log_id = str(uuid.uuid4()) visit_id = None conn = self._conn() try: with conn.cursor() as cur: # 1. Kiểm tra visit trong 30 phút gần nhất cur.execute( """ SELECT id FROM visits WHERE person_id = %s AND last_seen_at >= NOW() - INTERVAL '30 minutes' ORDER BY last_seen_at DESC LIMIT 1; """, (person_id,), ) row = cur.fetchone() if row: visit_id = str(row[0]) cur.execute( """ UPDATE visits SET last_seen_at = NOW(), appearance_count = appearance_count + 1 WHERE id = %s; """, (visit_id,), ) else: visit_id = str(uuid.uuid4()) cur.execute( """ INSERT INTO visits (id, person_id, first_seen_at, last_seen_at, appearance_count, created_at) VALUES (%s, %s, NOW(), NOW(), 1, NOW()); """, (visit_id, person_id), ) # 2. Ghi recognition_log đầy đủ field theo ERD cur.execute( """ INSERT INTO recognition_logs ( id, person_id, visit_id, embedding_id, raw_image_mongo_id, face_crop_mongo_id, recognition_status, similarity_distance, confidence, bbox_x, bbox_y, bbox_width, bbox_height, ui_color, ui_track_index, detected_at ) VALUES ( %s, %s, %s, %s, %s, %s, 'recognized', %s, %s, %s, %s, %s, %s, %s, %s, NOW() ); """, ( log_id, person_id, visit_id, embedding_id, raw_image_mongo_id, face_crop_mongo_id, float(1.0 - similarity), float(similarity), float(bbox.get("x1", 0)), float(bbox.get("y1", 0)), float(bbox.get("x2", 0) - bbox.get("x1", 0)), float(bbox.get("y2", 0) - bbox.get("y1", 0)), ui_color, ui_track_index, ), ) # 3. Cập nhật persons cur.execute( """ UPDATE persons SET total_appearances = total_appearances + 1, total_visits = (SELECT COUNT(*) FROM visits WHERE person_id = %s), last_seen_at = NOW(), updated_at = NOW() WHERE id = %s; """, (person_id, person_id), ) conn.commit() logger.info(f"[PG] Matched visit recorded: person={person_id}, visit={visit_id}") return visit_id, log_id except Exception as e: conn.rollback() logger.error(f"[PG] record_matched_visit failed: {e}") raise finally: conn.close() # ── record_unknown_log ─────────────────────────────────────────────── def record_unknown_log( self, face_crop_mongo_id: str, raw_image_mongo_id: str, similarity: float, bbox: dict, ui_color: str = "#FF5500", ui_track_index: int = 0, ) -> str: """ Ghi recognition_log với person_id = NULL cho khuôn mặt không nhận diện. Returns: log_id (str) """ log_id = str(uuid.uuid4()) conn = self._conn() try: with conn.cursor() as cur: cur.execute( """ INSERT INTO recognition_logs ( id, person_id, visit_id, embedding_id, raw_image_mongo_id, face_crop_mongo_id, recognition_status, similarity_distance, confidence, bbox_x, bbox_y, bbox_width, bbox_height, ui_color, ui_track_index, detected_at ) VALUES ( %s, NULL, NULL, NULL, %s, %s, 'unknown', %s, %s, %s, %s, %s, %s, %s, %s, NOW() ); """, ( log_id, raw_image_mongo_id, face_crop_mongo_id, float(1.0 - similarity), float(similarity), float(bbox.get("x1", 0)), float(bbox.get("y1", 0)), float(bbox.get("x2", 0) - bbox.get("x1", 0)), float(bbox.get("y2", 0) - bbox.get("y1", 0)), ui_color, ui_track_index, ), ) conn.commit() logger.info(f"[PG] Unknown face logged: log_id={log_id}") return log_id except Exception as e: conn.rollback() logger.error(f"[PG] record_unknown_log failed: {e}") raise finally: conn.close() # ── update_person_info ─────────────────────────────────────────────── def update_person_info( self, person_id: str, full_name: str, date_of_birth: str | None = None, phone_number: str | None = None, address: str | None = None, ) -> bool: """Cập nhật thông tin cá nhân của một person.""" conn = self._conn() try: with conn.cursor() as cur: cur.execute( """ UPDATE persons SET full_name = %s, date_of_birth = %s, phone_number = %s, address = %s, status = 'active', updated_at = NOW() WHERE id = %s; """, (full_name, date_of_birth, phone_number, address, person_id), ) updated = cur.rowcount > 0 conn.commit() return updated except Exception as e: conn.rollback() logger.error(f"[PG] update_person_info failed: {e}") raise finally: conn.close() # ── delete_person ──────────────────────────────────────────────────── def delete_person(self, person_id: str) -> tuple[bool, list[str]]: """ Xóa person và toàn bộ dữ liệu liên quan trên Supabase. Returns: (success, list of face_crop_mongo_ids cần xóa trên MongoDB) """ conn = self._conn() try: with conn.cursor() as cur: cur.execute("SELECT id FROM persons WHERE id = %s;", (person_id,)) if not cur.fetchone(): return False, [] # Thu thập tất cả mongo_ids cần xóa mongo_ids: list[str] = [] cur.execute( "SELECT face_crop_mongo_id FROM face_embeddings WHERE person_id = %s;", (person_id,), ) mongo_ids.extend([r[0] for r in cur.fetchall() if r[0]]) cur.execute( "SELECT face_crop_mongo_id FROM recognition_logs WHERE person_id = %s;", (person_id,), ) mongo_ids.extend([r[0] for r in cur.fetchall() if r[0]]) # Xóa theo thứ tự FK cur.execute("DELETE FROM recognition_logs WHERE person_id = %s;", (person_id,)) cur.execute("DELETE FROM visits WHERE person_id = %s;", (person_id,)) cur.execute("DELETE FROM face_embeddings WHERE person_id = %s;", (person_id,)) cur.execute("DELETE FROM persons WHERE id = %s;", (person_id,)) conn.commit() logger.info(f"[PG] Deleted person: {person_id}") return True, list(set(mongo_ids)) except Exception as e: conn.rollback() logger.error(f"[PG] delete_person failed: {e}") raise finally: conn.close() # ── delete_person_by_name ──────────────────────────────────────────── def delete_person_by_name(self, name: str) -> tuple[bool, list[str]]: """Xóa person theo full_name. Convenience wrapper.""" conn = self._conn() try: with conn.cursor() as cur: cur.execute("SELECT id FROM persons WHERE full_name = %s;", (name,)) row = cur.fetchone() if not row: return False, [] person_id = str(row[0]) return self.delete_person(person_id) finally: conn.close() # ── list_persons ───────────────────────────────────────────────────── def list_persons_with_faces(self) -> list[dict]: """ Lấy danh sách tất cả persons có face_embedding, kèm mongo_id ảnh crop. """ query = """ SELECT p.id::text AS person_id, p.full_name, p.status, p.total_appearances, p.total_visits, p.last_seen_at, e.id::text AS embedding_id, e.face_crop_mongo_id FROM persons p JOIN face_embeddings e ON p.id = e.person_id WHERE e.is_active = TRUE ORDER BY p.full_name; """ conn = self._conn() try: with conn.cursor(cursor_factory=RealDictCursor) as cur: cur.execute(query) rows = cur.fetchall() result = [] for r in rows: d = dict(r) if d.get("last_seen_at"): d["last_seen_at"] = d["last_seen_at"].isoformat() result.append(d) return result finally: conn.close() # ── Singleton ───────────────────────────────────────────────────────────── pg_storage = PostgresStorage()