AlexeerCT
2cd6b165d7
- Introduced a new validation function for SQLAlchemy URLs to ensure proper formatting and provide clearer error messages. - Updated the `get_database_url` function to include checks for both `TEST_DATABASE_URL` and `DATABASE_URL`, improving flexibility in configuration. - Removed the previous URL scheme validation in favor of the new validation method, streamlining the error handling process. These changes aim to improve the robustness and clarity of database connection handling in the application.
340 lines
12 KiB
Python
340 lines
12 KiB
Python
import importlib.util
|
|
import logging
|
|
import os
|
|
import sys
|
|
from logging.config import fileConfig
|
|
from urllib.parse import quote_plus
|
|
|
|
from alembic import context
|
|
from alembic.operations import ops
|
|
from core.config import settings
|
|
from core.database import Base
|
|
from sqlalchemy import engine_from_config, pool
|
|
from sqlalchemy.engine.url import make_url
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# this is the Alembic Config object, which provides
|
|
# access to the values within the .ini file in use.
|
|
config = context.config
|
|
|
|
|
|
def _strip_env_url(raw: str) -> str:
|
|
"""Quita espacios/comillas típicos de secretos CI (.env, Gitea)."""
|
|
url = raw.strip().strip('"').strip("'")
|
|
return url
|
|
|
|
|
|
def _normalize_alembic_sqlalchemy_url(url: str) -> str:
|
|
"""Alembic usa el driver síncrono psycopg2; normaliza DSN típicos de app/tests."""
|
|
url = _strip_env_url(url)
|
|
if url.startswith("postgresql+asyncpg://"):
|
|
return url.replace("postgresql+asyncpg://", "postgresql+psycopg2://", 1)
|
|
if url.startswith("postgresql+psycopg2://"):
|
|
return url
|
|
if url.startswith("postgresql://"):
|
|
return url.replace("postgresql://", "postgresql+psycopg2://", 1)
|
|
if url.startswith("postgres://"):
|
|
return url.replace("postgres://", "postgresql+psycopg2://", 1)
|
|
return url
|
|
|
|
|
|
def _validate_sqlalchemy_url(url: str, env_key: str) -> None:
|
|
"""Misma validación que create_engine; evita urlparse (falla con esquemas tipo postgresql+psycopg2)."""
|
|
try:
|
|
make_url(url)
|
|
except Exception as e:
|
|
raise RuntimeError(
|
|
f"{env_key} no es una URL de SQLAlchemy válida. "
|
|
"Ejemplo: postgresql://usuario:clave@host:5432/nombre_bd"
|
|
) from e
|
|
|
|
|
|
def get_database_url():
|
|
"""Obtiene la URL de la base de datos (PostgreSQL) desde variables de entorno o alembic.ini."""
|
|
# CI / pytest: misma URL que los tests (secret TEST_DATABASE_URL) o DATABASE_URL explícita.
|
|
# CRÍTICO: debe ser tupla con coma final si un solo elemento: ("X",) — si no, ("X") es str y el for
|
|
# itera caracteres; env_key "_" + os.environ["_"] (común en shells) rompe con URL inválida.
|
|
for env_key in ("TEST_DATABASE_URL", "DATABASE_URL"):
|
|
raw = os.environ.get(env_key)
|
|
if raw and raw.strip():
|
|
normalized = _normalize_alembic_sqlalchemy_url(raw)
|
|
_validate_sqlalchemy_url(normalized, env_key)
|
|
return normalized
|
|
|
|
# Construcción desde settings (CORE_DB_* en .env / entorno)
|
|
host = settings.CORE_DB_HOST
|
|
db = settings.CORE_DB_NAME
|
|
user = settings.CORE_DB_USER
|
|
password = settings.CORE_DB_PASSWORD
|
|
port = settings.CORE_DB_PORT
|
|
|
|
if host and db and user and password:
|
|
try:
|
|
encoded_user = quote_plus(user)
|
|
encoded_password = quote_plus(password)
|
|
encoded_db = quote_plus(db)
|
|
return f"postgresql+psycopg2://{encoded_user}:{encoded_password}@{host}:{port}/{encoded_db}"
|
|
except Exception as e:
|
|
logger.error(f"Error al construir URL: {e}")
|
|
|
|
# Fallback al archivo de configuración
|
|
url = config.get_main_option("sqlalchemy.url")
|
|
if not url or "${" in url or "%(" in url:
|
|
raise RuntimeError(
|
|
"No se ha configurado la cadena de conexión a PostgreSQL. "
|
|
"Define TEST_DATABASE_URL o DATABASE_URL, o variables CORE_DB_*; "
|
|
"sqlalchemy.url en alembic.ini con placeholders ${...} no está soportado."
|
|
)
|
|
|
|
return url
|
|
|
|
|
|
# Configurar la URL de la base de datos
|
|
database_url = get_database_url()
|
|
|
|
# Debug: mostrar la URL (sin la contraseña)
|
|
if os.environ.get("ALEMBIC_DEBUG"):
|
|
# Ocultar la contraseña para el debug en la URL
|
|
try:
|
|
before, after = database_url.split("@", 1)
|
|
if ":" in before:
|
|
before = before.split(":", 1)[0] + ":***"
|
|
debug_url = before + "@" + after
|
|
except Exception:
|
|
debug_url = "postgresql://***:***@***"
|
|
logger.error("Error al ocultar la contraseña en la URL para debug.")
|
|
|
|
config.set_main_option("sqlalchemy.url", database_url)
|
|
|
|
# Interpret the config file for Python logging.
|
|
# This line sets up loggers basically.
|
|
if config.config_file_name is not None:
|
|
fileConfig(config.config_file_name)
|
|
|
|
# Ajusta la ruta para que puedas importar core y módulos
|
|
BASE_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
|
sys.path.insert(0, BASE_DIR)
|
|
|
|
# Configuración de Alembic
|
|
config = context.config
|
|
fileConfig(config.config_file_name)
|
|
target_metadata = Base.metadata
|
|
|
|
|
|
def include_object(object_, name, type_, reflected, compare_to):
|
|
"""
|
|
Keep all objects in autogenerate.
|
|
FK noise is cleaned in process_revision_directives.
|
|
"""
|
|
return True
|
|
|
|
|
|
def _fk_drop_signature(op_):
|
|
if not isinstance(op_, ops.DropConstraintOp):
|
|
return None
|
|
if getattr(op_, "constraint_type", None) != "foreignkey":
|
|
return None
|
|
return (
|
|
getattr(op_, "schema", None),
|
|
getattr(op_, "table_name", None),
|
|
getattr(op_, "constraint_name", None),
|
|
)
|
|
|
|
|
|
def _fk_create_signature(op_):
|
|
if not isinstance(op_, ops.CreateForeignKeyOp):
|
|
return None
|
|
local_cols = tuple(getattr(op_, "local_cols", ()) or ())
|
|
remote_cols = tuple(getattr(op_, "remote_cols", ()) or ())
|
|
return (
|
|
getattr(op_, "source_schema", None),
|
|
getattr(op_, "source_table", None),
|
|
getattr(op_, "referent_schema", None),
|
|
getattr(op_, "referent_table", None),
|
|
local_cols,
|
|
remote_cols,
|
|
)
|
|
|
|
|
|
def _drop_to_create_match(drop_op, create_op):
|
|
if not isinstance(drop_op, ops.DropConstraintOp):
|
|
return False
|
|
if not isinstance(create_op, ops.CreateForeignKeyOp):
|
|
return False
|
|
if getattr(drop_op, "constraint_type", None) != "foreignkey":
|
|
return False
|
|
|
|
def _normalize_schema(value):
|
|
# PostgreSQL reports default schema inconsistently as None/public.
|
|
return "public" if value in (None, "") else value
|
|
|
|
# Prefer structural comparison using Alembic's reverse op when available.
|
|
reverse_create = getattr(drop_op, "_reverse", None)
|
|
if isinstance(reverse_create, ops.CreateForeignKeyOp):
|
|
return (
|
|
_normalize_schema(getattr(reverse_create, "source_schema", None))
|
|
== _normalize_schema(getattr(create_op, "source_schema", None))
|
|
and getattr(reverse_create, "source_table", None) == getattr(create_op, "source_table", None)
|
|
and _normalize_schema(getattr(reverse_create, "referent_schema", None))
|
|
== _normalize_schema(getattr(create_op, "referent_schema", None))
|
|
and getattr(reverse_create, "referent_table", None) == getattr(create_op, "referent_table", None)
|
|
and tuple(getattr(reverse_create, "local_cols", ()) or ())
|
|
== tuple(getattr(create_op, "local_cols", ()) or ())
|
|
and tuple(getattr(reverse_create, "remote_cols", ()) or ())
|
|
== tuple(getattr(create_op, "remote_cols", ()) or ())
|
|
)
|
|
|
|
# Fallback for older op payloads: compare source table/schema and name.
|
|
return (
|
|
_normalize_schema(getattr(drop_op, "schema", None)) == _normalize_schema(getattr(create_op, "source_schema", None))
|
|
and getattr(drop_op, "table_name", None) == getattr(create_op, "source_table", None)
|
|
and getattr(drop_op, "constraint_name", None) == getattr(create_op, "constraint_name", None)
|
|
)
|
|
|
|
|
|
def _prune_fk_churn(container):
|
|
if not hasattr(container, "ops"):
|
|
return
|
|
|
|
# First recurse into nested containers.
|
|
for op_ in list(container.ops):
|
|
_prune_fk_churn(op_)
|
|
|
|
table_ops = container.ops
|
|
kept_ops = []
|
|
consumed_indexes = set()
|
|
|
|
for i, op_i in enumerate(table_ops):
|
|
if i in consumed_indexes:
|
|
continue
|
|
|
|
if isinstance(op_i, ops.DropConstraintOp) and getattr(op_i, "constraint_type", None) == "foreignkey":
|
|
matched_j = None
|
|
for j in range(i + 1, len(table_ops)):
|
|
if j in consumed_indexes:
|
|
continue
|
|
op_j = table_ops[j]
|
|
if _drop_to_create_match(op_i, op_j):
|
|
matched_j = j
|
|
break
|
|
if matched_j is not None:
|
|
# Drop + recreate same FK detected; remove both.
|
|
consumed_indexes.add(i)
|
|
consumed_indexes.add(matched_j)
|
|
continue
|
|
|
|
kept_ops.append(op_i)
|
|
|
|
container.ops = kept_ops
|
|
|
|
|
|
def process_revision_directives(context_, revision, directives):
|
|
"""
|
|
Remove autogenerate noise where Alembic emits drop/create for equivalent FKs.
|
|
Real FK changes are preserved.
|
|
"""
|
|
if not directives:
|
|
return
|
|
script = directives[0]
|
|
_prune_fk_churn(script.upgrade_ops)
|
|
_prune_fk_churn(script.downgrade_ops)
|
|
|
|
|
|
def import_models_from_dir(dir_path: str):
|
|
"""Importa recursivamente cualquier archivo models.py desde dir_path y archivos en directorios models/"""
|
|
import sys
|
|
|
|
def _load_module(module_path: str):
|
|
rel_path = os.path.relpath(module_path, BASE_DIR)
|
|
module_name = rel_path.replace(os.sep, ".").replace(".py", "")
|
|
# Skip if already loaded to avoid duplicate SQLAlchemy table registrations
|
|
if module_name in sys.modules:
|
|
return
|
|
spec = importlib.util.spec_from_file_location(module_name, module_path)
|
|
mod = importlib.util.module_from_spec(spec)
|
|
# Register in sys.modules before exec so transitive imports resolve correctly
|
|
sys.modules[module_name] = mod
|
|
spec.loader.exec_module(mod)
|
|
|
|
for root, dirs, files in os.walk(dir_path):
|
|
# Importar archivos models.py directos
|
|
if "models.py" in files:
|
|
try:
|
|
_load_module(os.path.join(root, "models.py"))
|
|
except Exception as e:
|
|
logger.warning(f"No se pudo importar {os.path.join(root, 'models.py')}: {e}")
|
|
|
|
# Importar todos los archivos .py en directorios llamados "models"
|
|
if os.path.basename(root) == "models":
|
|
for file in files:
|
|
if file.endswith(".py") and not file.startswith("__"):
|
|
try:
|
|
_load_module(os.path.join(root, file))
|
|
except Exception as e:
|
|
logger.warning(f"No se pudo importar {os.path.join(root, file)}: {e}")
|
|
|
|
|
|
# Importar todos los models dentro de api/v1/modules y api/v1/modules/uploads
|
|
modules_dir = os.path.join(BASE_DIR, "api", "v1", "modules")
|
|
import_models_from_dir(modules_dir)
|
|
|
|
|
|
def run_migrations_offline() -> None:
|
|
"""Run migrations in 'offline' mode.
|
|
|
|
This configures the context with just a URL
|
|
and not an Engine, though an Engine is acceptable
|
|
here as well. By skipping the Engine creation
|
|
we don't even need a DBAPI to be available.
|
|
|
|
Calls to context.execute() here emit the given string to the
|
|
script output.
|
|
"""
|
|
url = config.get_main_option("sqlalchemy.url")
|
|
context.configure(
|
|
url=url,
|
|
target_metadata=target_metadata,
|
|
compare_type=True,
|
|
include_schemas=True,
|
|
include_object=include_object,
|
|
process_revision_directives=process_revision_directives,
|
|
literal_binds=True,
|
|
dialect_opts={"paramstyle": "named"},
|
|
)
|
|
|
|
with context.begin_transaction():
|
|
context.run_migrations()
|
|
|
|
|
|
def run_migrations_online() -> None:
|
|
"""Run migrations in 'online' mode.
|
|
|
|
In this scenario we need to create an Engine
|
|
and associate a connection with the context.
|
|
"""
|
|
connectable = engine_from_config(
|
|
config.get_section(config.config_ini_section, {}),
|
|
prefix="sqlalchemy.",
|
|
poolclass=pool.NullPool,
|
|
)
|
|
|
|
with connectable.connect() as connection:
|
|
context.configure(
|
|
connection=connection,
|
|
target_metadata=target_metadata,
|
|
compare_type=True,
|
|
include_schemas=True,
|
|
include_object=include_object,
|
|
process_revision_directives=process_revision_directives,
|
|
)
|
|
|
|
with context.begin_transaction():
|
|
context.run_migrations()
|
|
|
|
|
|
if context.is_offline_mode():
|
|
run_migrations_offline()
|
|
else:
|
|
run_migrations_online()
|