db_router.py

# -*- coding: utf-8 -*-
"""
db_router.py — Roteia Engine/SessionLocal para 'prod' (Load.db), 'test' (Load_teste.db)
e 'treinamento' (Load_treinamento.db), conforme escolha do usuário (mantida em st.session_state).

• set_db_choice("prod"|"test"|"treinamento")  → define o banco ativo para a sessão do usuário
• current_db_choice()                         → retorna 'prod' | 'test' | 'treinamento'
• get_engine()                                → engine do banco ativo (SQLite)
• get_session_factory()                       → sessionmaker do banco ativo
• SessionLocal()                              → sessão pronta no banco ativo

Observação:
- Este arquivo monta os caminhos dos bancos SQLite com base na pasta do projeto
  (onde está este arquivo), sem depender de .env.
- Se quiser usar .env futuramente, há uma seção comentada para isso.
"""

import os
import streamlit as st
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

# ============================
# Caminhos dos bancos (SQLite)
# ============================
# Pasta base do projeto (onde está este db_router.py)
BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# Nomes dos arquivos de banco conforme sua especificação
PROD_DB_NAME = "Load.db"
TEST_DB_NAME = "Load_teste.db"
TREINAMENTO_DB_NAME = "Load_treinamento.db"

# Monta URLs SQLite (formato: sqlite:///C:/.../arquivo.db)
DB1_PROD_URL = f"sqlite:///{os.path.join(BASE_DIR, PROD_DB_NAME)}"
DB2_TEST_URL = f"sqlite:///{os.path.join(BASE_DIR, TEST_DB_NAME)}"
DB3_TREINAMENTO_URL = f"sqlite:///{os.path.join(BASE_DIR, TREINAMENTO_DB_NAME)}"

# ============================
# (Opcional) Uso de .env
# ============================
# Se preferir usar .env futuramente, descomente abaixo e defina:
#   DB1_PROD_URL=sqlite:///C:/.../Load.db
#   DB2_TEST_URL=sqlite:///C:/.../Load_teste.db
#   DB3_TREINAMENTO_URL=sqlite:///C:/.../Load_treinamento.db
#
# from dotenv import load_dotenv
# load_dotenv()
# DB1_PROD_URL = os.getenv("DB1_PROD_URL", DB1_PROD_URL)
# DB2_TEST_URL = os.getenv("DB2_TEST_URL", DB2_TEST_URL)
# DB3_TREINAMENTO_URL = os.getenv("DB3_TREINAMENTO_URL", DB3_TREINAMENTO_URL)

# ============================
# Catálogo (helpers de UI — opcional)
# ============================
DB_URLS = {
    "prod":        DB1_PROD_URL,
    "test":        DB2_TEST_URL,
    "treinamento": DB3_TREINAMENTO_URL,
}

DB_LABELS = {
    "prod":        "Banco 1 (Produção)",
    "test":        "Banco 2 (Teste)",
    "treinamento": "Banco 3 (Treinamento)",
}

def list_banks() -> list[str]:
    """Lista as chaves de bancos disponíveis (opcional para UI)."""
    return list(DB_URLS.keys())

def bank_label(choice: str) -> str:
    """Rótulo amigável para a UI (opcional)."""
    return DB_LABELS.get(choice, choice)

# ============================
# Chaves de sessão
# ============================
SESSION_DB_CHOICE_KEY  = "__db_choice__"          # "prod" | "test" | "treinamento"
SESSION_DB_ENGINE_KEY  = "__db_engine__"          # cache de engine por escolha
SESSION_DB_FACTORY_KEY = "__db_session_factory__" # cache de sessionmaker por escolha

def set_db_choice(choice: str):
    """
    Define o banco ativo para a sessão do usuário.
    choice ∈ {"prod", "test", "treinamento"}.
    """
    choice = (choice or "").strip().lower()
    if choice not in DB_URLS:
        raise ValueError(f"db_choice inválido. Use uma destas chaves: {list(DB_URLS.keys())}")
    st.session_state[SESSION_DB_CHOICE_KEY] = choice
    # Ao trocar banco, invalida caches locais
    st.session_state.pop(SESSION_DB_ENGINE_KEY, None)
    st.session_state.pop(SESSION_DB_FACTORY_KEY, None)

def current_db_choice() -> str:
    """Retorna 'prod' | 'test' | 'treinamento'. Default: 'prod'."""
    return st.session_state.get(SESSION_DB_CHOICE_KEY, "prod")

def _url_for_choice(choice: str) -> str:
    return DB_URLS[choice]

def get_engine():
    """
    Entrega o engine do banco ATIVO (por sessão). Cria se não existir.
    """
    choice = current_db_choice()
    cached = st.session_state.get(SESSION_DB_ENGINE_KEY)
    if cached and getattr(cached, "__db_choice__", None) == choice:
        return cached

    url = _url_for_choice(choice)
    engine_args = {
        "echo": False,
        "pool_pre_ping": True,
        # check_same_thread=False para permitir uso em múltiplas threads do Streamlit
        "connect_args": {"check_same_thread": False} if url.startswith("sqlite") else {},
    }

    eng = create_engine(url, **engine_args)
    setattr(eng, "__db_choice__", choice)
    st.session_state[SESSION_DB_ENGINE_KEY] = eng
    return eng

def get_session_factory():
    """
    Entrega um sessionmaker vinculado ao engine do banco ATIVO.
    """
    choice = current_db_choice()
    fac = st.session_state.get(SESSION_DB_FACTORY_KEY)
    if fac and getattr(fac, "__db_choice__", None) == choice:
        return fac

    fac = sessionmaker(bind=get_engine(), autocommit=False, autoflush=False)
    setattr(fac, "__db_choice__", choice)
    st.session_state[SESSION_DB_FACTORY_KEY] = fac
    return fac

def SessionLocal():
    """
    Cria uma sessão (sempre no banco ativo).
    Uso:
        db = SessionLocal()
        try:
            ...
        finally:
            db.close()
    """
    return get_session_factory()()