Update db_router.py

db_router.py

@@ -1,152 +1,220 @@
-
-# -*- 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()()
-
-
+# -*- coding: utf-8 -*-
+"""
+db_router.py — Roteia Engine/SessionLocal para:
+ - 'prod'         → Load.db
+ - 'test'         → Load_teste.db
+ - 'treinamento'  → Load_treinamento.db
+
+Mantém a escolha do usuário em st.session_state (quando disponível).
+Exibe uma API simples esperada pelo banco.py:
+
+• set_db_choice("prod"|"test"|"treinamento")  → define o banco ativo
+• current_db_choice()                         → retorna 'prod' | 'test' | 'treinamento'
+• get_engine()                                → engine do banco ativo (SQLite)
+• get_session_factory()                       → sessionmaker do banco ativo
+• SessionLocal()                              → sessão no banco ativo
+
+Compatibilidade:
+- Se Streamlit não estiver disponível (execução fora do app), usa cache global interno (fallback).
+- Pode ler URLs do .env/Secrets (opcional) para apontar bancos em outro local/driver.
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Dict, Optional
+
+# Streamlit é preferível; mas se não houver (execução fora do app), caímos em fallback
+try:
+    import streamlit as st
+    _HAS_ST = True
+except Exception:
+    _HAS_ST = False
+
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+
+# ============================
+# Configuração e caminhos base
+# ============================
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+# Nomes de arquivos de banco conforme sua especificação (com caixa)
+PROD_DB_NAME         = "Load.db"
+TEST_DB_NAME         = "Load_teste.db"
+TREINAMENTO_DB_NAME  = "Load_treinamento.db"
+
+# ============================
+# (Opcional) uso de .env/Secrets
+# ============================
+# Se quiser apontar para outros caminhos/URIs via Secrets do Spaces:
+#   DB1_PROD_URL, DB2_TEST_URL, DB3_TREINAMENTO_URL
+# Se não definir, usamos o SQLite no diretório do projeto.
+DB1_PROD_URL = os.getenv("DB1_PROD_URL", f"sqlite:///{os.path.join(BASE_DIR, PROD_DB_NAME)}")
+DB2_TEST_URL = os.getenv("DB2_TEST_URL", f"sqlite:///{os.path.join(BASE_DIR, TEST_DB_NAME)}")
+DB3_TREINAMENTO_URL = os.getenv("DB3_TREINAMENTO_URL", f"sqlite:///{os.path.join(BASE_DIR, TREINAMENTO_DB_NAME)}")
+
+DB_URLS: Dict[str, str] = {
+    "prod":        DB1_PROD_URL,
+    "test":        DB2_TEST_URL,
+    "treinamento": DB3_TREINAMENTO_URL,
+}
+
+DB_LABELS: Dict[str, str] = {
+    "prod":        "Banco 1 (Produção)",
+    "test":        "Banco 2 (Teste)",
+    "treinamento": "Banco 3 (Treinamento)",
+}
+
+# ============================
+# (Opcional) normalização de case para SQLite
+# ============================
+# Se você quiser garantir que exista também uma cópia "load.db" (minúsculo)
+# a partir de "Load.db" (maiúsculo) — útil se outros módulos esperarem minúsculas:
+_ENABLE_SQLITE_LOWERCASE_ALIAS = False  # mude para True se quiser habilitar
+
+def _ensure_sqlite_lowercase_alias():
+    import shutil
+    # Para cada caminho SQLite "Load*.db", garantir um "load*.db"
+    for key, url in DB_URLS.items():
+        if not url.startswith("sqlite:///"):
+            continue
+        path = url.replace("sqlite:///", "", 1)
+        # Se usar caminho relativo da BASE_DIR, resolva:
+        if not os.path.isabs(path):
+            path = os.path.join(BASE_DIR, path)
+        dirn, fname = os.path.split(path)
+        lower_name = fname.lower()
+        lower_path = os.path.join(dirn, lower_name)
+        if fname != lower_name and os.path.exists(path) and not os.path.exists(lower_path):
+            try:
+                shutil.copy(path, lower_path)
+            except Exception:
+                # Se falhar, ignora (não é crítico)
+                pass
+
+if _ENABLE_SQLITE_LOWERCASE_ALIAS:
+    _ensure_sqlite_lowercase_alias()
+
+# ============================
+# Helpers de UI
+# ============================
+def list_banks() -> list[str]:
+    """Lista as chaves de bancos disponíveis (para UI)."""
+    return list(DB_URLS.keys())
+
+def bank_label(choice: str) -> str:
+    """Rótulo amigável para a UI."""
+    return DB_LABELS.get(choice, choice)
+
+# ============================
+# Chaves de sessão (ou fallback global)
+# ============================
+SESSION_DB_CHOICE_KEY  = "__db_choice__"          # "prod" | "test" | "treinamento"
+SESSION_DB_ENGINE_KEY  = "__db_engine__"          # cache de engine
+SESSION_DB_FACTORY_KEY = "__db_session_factory__" # cache de sessionmaker
+
+# Fallback global quando não houver Streamlit
+_GLOBAL_STATE: Dict[str, object] = {
+    SESSION_DB_CHOICE_KEY: "prod",
+    SESSION_DB_ENGINE_KEY: None,
+    SESSION_DB_FACTORY_KEY: None,
+}
+
+def _session_get(key: str, default=None):
+    if _HAS_ST:
+        return st.session_state.get(key, default)
+    return _GLOBAL_STATE.get(key, default)
+
+def _session_set(key: str, value):
+    if _HAS_ST:
+        st.session_state[key] = value
+    else:
+        _GLOBAL_STATE[key] = value
+
+def _session_pop(key: str):
+    if _HAS_ST:
+        st.session_state.pop(key, None)
+    else:
+        _GLOBAL_STATE.pop(key, None)
+
+# ============================
+# Escolha do banco
+# ============================
+def set_db_choice(choice: str):
+    """
+    Define o banco ativo para a sessão atual.
+    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())}")
+    _session_set(SESSION_DB_CHOICE_KEY, choice)
+    # Ao trocar, invalida caches
+    _session_pop(SESSION_DB_ENGINE_KEY)
+    _session_pop(SESSION_DB_FACTORY_KEY)
+
+def current_db_choice() -> str:
+    """Retorna 'prod' | 'test' | 'treinamento' (default: 'prod')."""
+    val = _session_get(SESSION_DB_CHOICE_KEY, "prod")
+    if val not in DB_URLS:
+        val = "prod"
+        _session_set(SESSION_DB_CHOICE_KEY, val)
+    return val
+
+# ============================
+# Engine / Session por ambiente
+# ============================
+def _url_for_choice(choice: str) -> str:
+    return DB_URLS[choice]
+
+def _engine_args_for_url(url: str) -> dict:
+    args = {
+        "echo": False,
+        "pool_pre_ping": True,
+    }
+    if url.startswith("sqlite:///"):
+        # evita erro em threads do Streamlit
+        args["connect_args"] = {"check_same_thread": False}
+    return args
+
+def get_engine():
+    """
+    Entrega o engine do banco ATIVO (por sessão). Cria e cacheia se necessário.
+    """
+    choice = current_db_choice()
+    cached = _session_get(SESSION_DB_ENGINE_KEY)
+    if cached is not None and getattr(cached, "__db_choice__", None) == choice:
+        return cached
+
+    url = _url_for_choice(choice)
+    eng = create_engine(url, **_engine_args_for_url(url))
+    setattr(eng, "__db_choice__", choice)
+    _session_set(SESSION_DB_ENGINE_KEY, eng)
+    return eng
+
+def get_session_factory():
+    """
+    Entrega um sessionmaker vinculado ao engine do banco ATIVO (em cache).
+    """
+    choice = current_db_choice()
+    fac = _session_get(SESSION_DB_FACTORY_KEY)
+    if fac is not None and getattr(fac, "__db_choice__", None) == choice:
+        return fac
+
+    fac = sessionmaker(bind=get_engine(), autocommit=False, autoflush=False)
+    setattr(fac, "__db_choice__", choice)
+    _session_set(SESSION_DB_FACTORY_KEY, fac)
+    return fac
+
+def SessionLocal():
+    """
+    Cria uma sessão no banco ATIVO.
+    Uso:
+        db = SessionLocal()
+        try:
+            ...
+        finally:
+            db.close()
+    """
+    return get_session_factory()()