Update db_router.py

db_router.py

@@ -5,22 +5,27 @@ db_router.py — Roteia Engine/SessionLocal para:
  - 'test'         → Load_teste.db
  - 'treinamento'  → Load_treinamento.db
 
-Mantém a escolha do usuário em st.session_state (quando disponível).
-Expõe:
-• set_db_choice("prod"|"test"|"treinamento")
-• current_db_choice()
-• get_engine()
-• get_session_factory()
-• SessionLocal()
-
-Inclui garantia de criação do diretório pai do SQLite.
+Mantém a escolha do usuário em st.session_state (quando disponível) e em variável
+de ambiente (DB_CHOICE) para persistir entre reruns/contexts.
+
+APIs expostas (compatíveis com o app):
+• get_available_choices() -> list[str]
+• set_current_db_choice(choice: str) -> None
+• set_db_choice(choice: str) -> None
+• current_db_choice() -> str
+• bank_label(choice: str) -> str
+• get_engine() -> sqlalchemy.Engine
+• get_session_factory() -> sqlalchemy.orm.sessionmaker
+• SessionLocal() -> sqlalchemy.orm.Session
+
+Inclui garantia de criação do diretório pai do SQLite com fallback para ~/.ioirun.
 Compatível com execução fora do Streamlit (fallback em estado global).
 """
 
 from __future__ import annotations
 
 import os
-from typing import Dict, Optional
+from typing import Dict, Optional, Any
 
 # Streamlit é preferível; mas se não houver (execução fora do app), caímos em fallback
 try:
@@ -37,7 +42,7 @@ from sqlalchemy.orm import sessionmaker
 # ============================
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 
-# Nomes de arquivos de banco conforme sua especificação (com caixa)
+# Nomes de arquivos de banco conforme sua especificação
 PROD_DB_NAME         = "Load.db"
 TEST_DB_NAME         = "Load_teste.db"
 TREINAMENTO_DB_NAME  = "Load_treinamento.db"
@@ -53,10 +58,15 @@ DB_URLS: Dict[str, str] = {
     "treinamento": DB3_TREINAMENTO_URL,
 }
 
+# Aliases aceitos (ex.: "train" → "treinamento")
+CHOICE_ALIASES: Dict[str, str] = {
+    "train": "treinamento",
+}
+
 DB_LABELS: Dict[str, str] = {
-    "prod":        "Banco 1 (Produção)",
-    "test":        "Banco 2 (Teste)",
-    "treinamento": "Banco 3 (Treinamento)",
+    "prod":        "🟢 Produção",
+    "test":        "🔴 Teste",
+    "treinamento": "🟡 Treinamento",
 }
 
 # ============================
@@ -69,15 +79,17 @@ def _ensure_parent_dir_sqlite(url: str) -> str:
     """
     if not url or not url.startswith("sqlite"):
         return url
-    path = url.replace("sqlite:///", "", 1)
-    if path.startswith("//"):
-        path = path[1:]
+    # Extrai caminho local do SQLite (sqlite:////abs/path  ou sqlite:///rel/path)
+    prefix = "sqlite:///"
+    path = url[len(prefix):] if url.startswith(prefix) else url
+    # Normaliza e cria parent
     file_path = os.path.abspath(path)
     parent = os.path.dirname(file_path)
     try:
         os.makedirs(parent, exist_ok=True)
         return url
     except Exception:
+        # Fallback para HOME gravável
         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))
@@ -86,10 +98,14 @@ def _ensure_parent_dir_sqlite(url: str) -> str:
 # ============================
 # Helpers de UI
 # ============================
-def list_banks() -> list[str]:
+def get_available_choices() -> list[str]:
     """Lista as chaves de bancos disponíveis (para UI)."""
     return list(DB_URLS.keys())
 
+def list_banks() -> list[str]:
+    """Compat anterior — alias de get_available_choices()."""
+    return get_available_choices()
+
 def bank_label(choice: str) -> str:
     """Rótulo amigável para a UI."""
     return DB_LABELS.get(choice, choice)
@@ -102,8 +118,8 @@ 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",
+_GLOBAL_STATE: Dict[str, Any] = {
+    SESSION_DB_CHOICE_KEY: os.getenv("DB_CHOICE", "prod"),
     SESSION_DB_ENGINE_KEY: None,
     SESSION_DB_FACTORY_KEY: None,
 }
@@ -125,28 +141,69 @@ def _session_pop(key: str):
     else:
         _GLOBAL_STATE.pop(key, None)
 
+# ============================
+# Normalização da escolha
+# ============================
+def _normalize_choice(raw: Optional[str]) -> str:
+    val = (raw or "").strip().lower()
+    if val in CHOICE_ALIASES:
+        val = CHOICE_ALIASES[val]
+    if val not in DB_URLS:
+        val = "prod"
+    return val
+
 # ============================
 # Escolha do banco
 # ============================
 def set_db_choice(choice: str):
     """
     Define o banco ativo para a sessão atual.
-    choice ∈ {"prod", "test", "treinamento"}.
+    choice ∈ {"prod", "test", "treinamento"} (aceita alias "train" → "treinamento").
+    Invalida caches de engine/session e atualiza ENV (DB_CHOICE).
     """
-    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())}")
+    choice = _normalize_choice(choice)
+
+    # Se houver engine cacheado, faça dispose antes de trocar
+    cached_engine = _session_get(SESSION_DB_ENGINE_KEY)
+    if cached_engine is not None:
+        try:
+            cached_engine.dispose()
+        except Exception:
+            pass
+
     _session_set(SESSION_DB_CHOICE_KEY, choice)
+    os.environ["DB_CHOICE"] = choice  # permite que outros módulos leiam
+
     # Ao trocar, invalida caches
     _session_pop(SESSION_DB_ENGINE_KEY)
     _session_pop(SESSION_DB_FACTORY_KEY)
 
+def set_current_db_choice(choice: str) -> None:
+    """Alias compatível com app: set_current_db_choice(choice)."""
+    set_db_choice(choice)
+
 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:
+    """
+    Retorna 'prod' | 'test' | 'treinamento' (default: 'prod').
+    Prioriza session_state; se ausente, lê DB_CHOICE do ambiente.
+    """
+    # 1) Session
+    val = _session_get(SESSION_DB_CHOICE_KEY)
+    # 2) ENV (permite seleção por URL/externo)
+    if val is None:
+        env_val = os.getenv("DB_CHOICE")
+        if env_val:
+            val = _normalize_choice(env_val)
+            _session_set(SESSION_DB_CHOICE_KEY, val)
+    # 3) Default
+    if val is None:
         val = "prod"
         _session_set(SESSION_DB_CHOICE_KEY, val)
+
+    # Sanitize
+    val = _normalize_choice(val)
+    if val != _session_get(SESSION_DB_CHOICE_KEY):
+        _session_set(SESSION_DB_CHOICE_KEY, val)
     return val
 
 # ============================
@@ -175,7 +232,7 @@ def get_engine():
         return cached
 
     url = _url_for_choice(choice)
-    url = _ensure_parent_dir_sqlite(url)  # ⬅️ garante diretório pai
+    url = _ensure_parent_dir_sqlite(url)  # ⬅️ garante diretório pai se for SQLite
 
     eng = create_engine(url, **_engine_args_for_url(url))
     setattr(eng, "__db_choice__", choice)