import os, sys, re from typing import List, Dict, Optional from sqlalchemy import text from sqlalchemy.exc import SQLAlchemyError # Agregamos el path para encontrar core/DBConnect.py # sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from app.core.DBConnect import DatabaseConfig, SessionManager class TablePartitioner: """ Clase genérica para manejo de tablas particionadas en PostgreSQL. """ def __init__(self, schema: str, table_name: str, session_manager: Optional[SessionManager] = None): self.schema = schema self.table_name = table_name if session_manager is None: db_config = DatabaseConfig() self.session_manager = SessionManager(db_config) else: self.session_manager = session_manager # ---------------------------------------------------------------------- # Helpers internos # ---------------------------------------------------------------------- def _generate_partition_name(self, rfc_hash_usr_ownr: str) -> str: """Genera nombre de partición seguro""" safe_key = re.sub(r'[^a-zA-Z0-9_]', '_', rfc_hash_usr_ownr.lower()) return f"{self.table_name}_{safe_key}" def _full_table_name(self) -> str: """Retorna nombre completo de la tabla""" return f"{self.schema}.{self.table_name}" def _partition_exists(self, session, partition_name: str) -> bool: """Verifica si una partición existe""" try: query = text(""" SELECT 1 FROM pg_tables WHERE schemaname = :schema AND tablename = :partition_name """) result = session.execute(query, { "schema": self.schema, "partition_name": partition_name }).fetchone() return result is not None except SQLAlchemyError: return False # ---------------------------------------------------------------------- # Creación de tabla maestra # ---------------------------------------------------------------------- def create_master_table_if_not_exists(self, columns_definition: str, force_recreate: bool = False) -> bool: """Crea la tabla maestra particionada""" session = self.session_manager.get_session() try: # Verificar si la tabla ya existe check_query = text(""" SELECT 1 FROM information_schema.tables WHERE table_schema = :schema AND table_name = :table_name """) exists = session.execute(check_query, { "schema": self.schema, "table_name": self.table_name }).fetchone() if exists and not force_recreate: print(f"✓ Tabla maestra {self._full_table_name()} ya existe") return True if force_recreate and exists: print(f"⚠️ Eliminando tabla existente {self._full_table_name()}") session.execute(text(f"DROP TABLE IF EXISTS {self._full_table_name()} CASCADE;")) # Crear tabla particionada query = f""" CREATE TABLE IF NOT EXISTS {self._full_table_name()} ( {columns_definition} ) PARTITION BY LIST (rfc_hash_usr_ownr); """ session.execute(text(query)) session.commit() print(f"✅ Tabla maestra creada: {self._full_table_name()}") return True except SQLAlchemyError as e: print(f"❌ Error creando tabla maestra {self._full_table_name()}: {e}") session.rollback() return False finally: self.session_manager.close_session(session) # ---------------------------------------------------------------------- # Manejo de particiones # ---------------------------------------------------------------------- def list_partitions(self) -> List[str]: """Lista todas las particiones de la tabla""" session = self.session_manager.get_session() try: query = text(""" SELECT schemaname || '.' || tablename AS partition_name FROM pg_tables WHERE schemaname = :schema AND tablename LIKE :pattern ORDER BY tablename; """) result = session.execute(query, { "schema": self.schema, "pattern": f"{self.table_name}_%" }) return [row[0] for row in result.fetchall()] except SQLAlchemyError as e: print(f"❌ Error listando particiones de {self._full_table_name()}: {e}") return [] finally: self.session_manager.close_session(session) def create_partition(self, rfc_hash_usr_ownr: str) -> bool: """Crea una partición para el rfc_hash_usr_ownr específico""" partition_name = self._generate_partition_name(rfc_hash_usr_ownr) session = self.session_manager.get_session() try: # Verificar si ya existe if self._partition_exists(session, partition_name): print(f"ℹ️ Partición ya existe: {partition_name}") return True # Crear la partición full_partition_name = f"{self.schema}.{partition_name}" query = f""" CREATE TABLE {full_partition_name} PARTITION OF {self._full_table_name()} FOR VALUES IN ('{rfc_hash_usr_ownr}'); """ session.execute(text(query)) session.commit() print(f"✅ Partición creada: {partition_name}") return True except SQLAlchemyError as e: print(f"❌ Error creando partición {partition_name}: {e}") session.rollback() return False finally: self.session_manager.close_session(session) def drop_partition(self, rfc_hash_usr_ownr: str) -> bool: """Elimina una partición específica""" partition_name = self._generate_partition_name(rfc_hash_usr_ownr) session = self.session_manager.get_session() try: full_partition_name = f"{self.schema}.{partition_name}" session.execute(text(f"DROP TABLE IF EXISTS {full_partition_name} CASCADE;")) session.commit() print(f"✅ Partición eliminada: {partition_name}") return True except SQLAlchemyError as e: print(f"❌ Error eliminando partición {partition_name}: {e}") session.rollback() return False finally: self.session_manager.close_session(session) # ---------------------------------------------------------------------- # Estadísticas y administración # ---------------------------------------------------------------------- def get_partition_stats(self) -> List[Dict]: """Obtiene estadísticas de todas las particiones""" session = self.session_manager.get_session() try: query = text(""" SELECT relname AS partition_name, n_live_tup AS row_estimate, pg_size_pretty(pg_total_relation_size(oid)) as size, pg_total_relation_size(oid) as size_bytes FROM pg_stat_user_tables t JOIN pg_class c ON c.relname = t.relname WHERE t.schemaname = :schema AND t.relname LIKE :table_pattern ORDER BY t.relname; """) result = session.execute(query, { "schema": self.schema, "table_pattern": f"{self.table_name}_%" }) return [ { "partition": row[0], "rows": row[1] or 0, "size": row[2], "size_bytes": row[3] or 0 } for row in result.fetchall() ] except SQLAlchemyError as e: print(f"❌ Error obteniendo estadísticas de particiones: {e}") return [] finally: self.session_manager.close_session(session) def optimize_partition(self, rfc_hash_usr_ownr: str) -> bool: """Optimiza la partición usando VACUUM y ANALYZE""" partition_name = self._generate_partition_name(rfc_hash_usr_ownr) session = self.session_manager.get_session() try: if not self._partition_exists(session, partition_name): print(f"⚠️ Partición no existe: {partition_name}") return False full_partition_name = f"{self.schema}.{partition_name}" # PostgreSQL requiere que VACUUM se ejecute fuera de transacción session.commit() # Commit cualquier transacción pendiente session.execute(text(f"VACUUM ANALYZE {full_partition_name};")) print(f"✅ Partición optimizada: {partition_name}") return True except SQLAlchemyError as e: print(f"❌ Error optimizando partición {partition_name}: {e}") return False finally: self.session_manager.close_session(session) def get_partition_info(self, rfc_hash_usr_ownr: str) -> Dict: """Obtiene información detallada de una partición específica""" partition_name = self._generate_partition_name(rfc_hash_usr_ownr) session = self.session_manager.get_session() try: if not self._partition_exists(session, partition_name): return {"exists": False} query = text(f""" SELECT COUNT(*) as record_count, pg_size_pretty(pg_total_relation_size('{self.schema}.{partition_name}')) as size, pg_total_relation_size('{self.schema}.{partition_name}') as size_bytes FROM {self.schema}.{partition_name}; """) result = session.execute(query).fetchone() return { "exists": True, "partition_name": partition_name, "full_name": f"{self.schema}.{partition_name}", "record_count": result[0] if result else 0, "size": result[1] if result else "0 bytes", "size_bytes": result[2] if result else 0, "rfc_hash_usr_ownr": rfc_hash_usr_ownr } except SQLAlchemyError as e: print(f"❌ Error obteniendo información de partición {partition_name}: {e}") return {"exists": False, "error": str(e)} finally: self.session_manager.close_session(session)