IOI-RUN / banco.py
Roudrigus's picture
Update banco.py
a7cdbad verified
raw
history blame
8.36 kB
# -*- coding: utf-8 -*-
"""
banco.py
Compatível com:
- Roteamento por ambiente (db_router.py): produção/teste/treinamento
- Fallback: um único DATABASE_URL vindo de env/Secrets
- Postgres / MySQL / SQLite (c/ alias de case para load.db no Linux)
- Garantia de criação do diretório pai do SQLite (evita 'unable to open database file')
"""
import os
import shutil
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, declarative_base
from dotenv import load_dotenv
# Caminho base do projeto
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
# Carrega variáveis (.env) — no Spaces você usa Settings → Secrets
load_dotenv()
# =========================================================
# 1) Correção de case para SQLite (Load.db → load.db) — opcional
# =========================================================
def _ensure_sqlite_case_alias() -> str:
"""
Garante que exista 'load.db' no diretório do app.
Se encontrar 'Load.db' (ou outra variação de caixa), copia para 'load.db'.
Retorna o caminho absoluto de 'load.db'.
"""
lower = os.path.join(BASE_DIR, "load.db")
if os.path.exists(lower):
return lower
# Candidatos com caixa diferente
for cand in ("Load.db", "LOAD.DB", "Load.DB"):
up = os.path.join(BASE_DIR, cand)
if os.path.exists(up):
try:
shutil.copy(up, lower)
except Exception:
# Se falhar a cópia, segue adiante; o sqlite criará um vazio no primeiro uso
pass
break
return lower
# =========================================================
# 2) Helper: garantir diretório pai do arquivo SQLite
# =========================================================
def _ensure_parent_dir_sqlite(sqlite_url: str) -> str:
"""
Garante que o diretório pai do arquivo SQLite exista.
Caso não consiga criar (permissão), cai para ~/.ioirun/<arquivo>.db
Retorna a URL (que pode ser ajustada para fallback).
"""
if not sqlite_url or not sqlite_url.startswith("sqlite"):
return sqlite_url
# Extrai caminho do arquivo a partir da URL
# Formatos: sqlite:///relativo.db | sqlite:////abs/path/to.db
path = sqlite_url.replace("sqlite:///", "", 1)
# se vier com // (pouco comum), normaliza
if path.startswith("//"):
path = path[1:]
file_path = os.path.abspath(path)
parent = os.path.dirname(file_path)
try:
os.makedirs(parent, exist_ok=True)
return sqlite_url
except Exception:
# fallback para HOME (gravável no Spaces)
home_dir = os.path.join(os.path.expanduser("~"), ".ioirun")
os.makedirs(home_dir, exist_ok=True)
alt = os.path.join(home_dir, os.path.basename(file_path))
return f"sqlite:///{alt}"
# =========================================================
# 3) Suporte a roteador (db_router.py) — opcional
# =========================================================
# Se existir um roteador, delegamos a criação da engine e da SessionLocal
# conforme o "banco atual" selecionado (ex.: prod/test/treinamento).
try:
from db_router import (
get_engine as _router_get_engine,
get_session_factory as _router_get_session_factory,
SessionLocal as _router_SessionLocal,
)
_HAS_ROUTER = True
except Exception:
_HAS_ROUTER = False
# =========================================================
# 4) Fallback quando NÃO há roteador: construir a URI
# =========================================================
def _build_fallback_uri() -> str:
"""
Monta a URI do banco quando não existe db_router.
Ordem de preferência:
1. DATABASE_URL (completo)
2. Variáveis separadas: DB_DRIVER, DB_HOST, DB_PORT, DB_USER, DB_PASS, DB_NAME
3. SQLite local em 'load.db'
"""
# 4.1 DATABASE_URL completo
url = os.getenv("DATABASE_URL")
if url:
# Garante diretório pai no caso de sqlite
return _ensure_parent_dir_sqlite(url)
# 4.2 Campos separados
driver = (os.getenv("DB_DRIVER") or "").strip().lower() # "postgresql", "mysql"
host = os.getenv("DB_HOST")
port = os.getenv("DB_PORT")
user = os.getenv("DB_USER")
pwd = os.getenv("DB_PASS")
name = os.getenv("DB_NAME")
if driver and host and user and pwd and name:
# Defaults de porta
if not port:
port = "5432" if driver.startswith("post") else "3306"
if driver.startswith("post"): # PostgreSQL
return f"postgresql+psycopg2://{user}:{pwd}@{host}:{port}/{name}"
elif driver.startswith("mysql"): # MySQL/MariaDB
return f"mysql+pymysql://{user}:{pwd}@{host}:{port}/{name}"
# 4.3 SQLite local (fallback)
sqlite_path = _ensure_sqlite_case_alias()
return _ensure_parent_dir_sqlite(f"sqlite:///{sqlite_path}")
# =========================================================
# 5) Engine / SessionLocal
# =========================================================
# Observação importante:
# - Se houver db_router, usamos as fábricas do roteador (engine/sessões por ambiente).
# - Caso contrário, criamos uma engine única a partir de DATABASE_URL/DB_* ou SQLite.
if _HAS_ROUTER:
# =========== Com roteador ===========
def get_engine():
"""
Retorna a engine do banco ATUAL (prod/test/treinamento),
conforme implementado pelo seu db_router.get_engine().
"""
return _router_get_engine()
def _session_factory():
"""
Fábrica de Session para o banco ATUAL (vinda do roteador).
"""
return _router_get_session_factory()
# SessionLocal fornecida pelo roteador (respeita o banco atual)
SessionLocal = _router_SessionLocal
else:
# =========== Fallback sem roteador ===========
DATABASE_URL = _build_fallback_uri()
engine_args = {
"echo": False, # defina True para depuração de SQL
"pool_pre_ping": True, # valida conexões antes de usar
}
if DATABASE_URL.startswith("sqlite"):
# Parâmetros específicos para SQLite
engine_args["connect_args"] = {"check_same_thread": False}
_engine = create_engine(DATABASE_URL, **engine_args)
def get_engine():
"""
Engine fixa do fallback. Se quiser trocar de banco em runtime sem roteador,
você precisará recriar a engine manualmente (ou adotar db_router).
"""
return _engine
_SessionFactory = sessionmaker(
autocommit=False,
autoflush=False,
bind=_engine,
)
def _session_factory():
return _SessionFactory
# Para compatibilidade com código que usa SessionLocal()
SessionLocal = _SessionFactory
# =========================================================
# 6) Expor 'engine' e Base ORM
# =========================================================
# Atenção: 'engine' é resolvido no momento da importação.
# Se você troca de banco após importar 'banco', prefira usar get_engine()
# e criar a sessão com SessionLocal() para assegurar o banco ATUAL.
engine = get_engine()
Base = declarative_base()
# =========================================================
# 7) Utilitários (opcionais)
# =========================================================
def init_schema():
"""
Cria/atualiza as tabelas no banco ATUAL.
• Com roteador: aplica no banco escolhido (Produção/Teste/...).
• Sem roteador: aplica no DATABASE_URL padrão.
Use em DEV/TESTE; em produção, prefira migrações (ex.: Alembic).
"""
# Importa 'models' de forma tardia/segura para registrar mapeamentos
import importlib
try:
importlib.import_module("models")
except ModuleNotFoundError:
# Ajuste se seus modelos estiverem em outro pacote
# importlib.import_module("app.models")
raise
Base.metadata.create_all(bind=get_engine())
def db_info() -> dict:
"""
Retorna informações básicas do banco ativo (para debug/UX).
"""
eng = get_engine()
try:
url = str(eng.url)
except Exception:
try:
url = DATABASE_URL # type: ignore[name-defined]
except Exception:
url = "(não disponível)"
return {
"url": url,
"using_router": _HAS_ROUTER,
}