ec26183e2e
- ruff.toml: select E/W/F/I/N rules, line-length=120, drop legacy ignores - Auto-fix: sort 82 import blocks (isort), remove 29 unused imports, strip 6 trailing-whitespace blank lines in docstrings - main.py: move setup_logging and settings imports to top (E402) - errors.py: noqa N818 on DDD exception names (96 call sites, safe) - intel_service.py: noqa N817 for universal ET alias - atomic/elastic/sigma import services: move _MAX_UNCOMPRESSED_SIZE and _MAX_ENTRIES to module level (N806) - compliance_import_service.py: move SAMPLE_CONTROLS / CIS_CONTROLS to module level; wrap long description strings (N806 + E501) - snapshot_service.py: move STATUS_ORDER dict to module level (N806) - sigma_import_service.py: remove dead dedup_key expression (F841) - threat_actor_import_service.py: remove dead stix_to_actor expression (F841) - data_source.py, seed_demo.py, campaign_scheduler_service.py, lolbas_import_service.py: wrap lines exceeding 120 chars (E501) - d3fend_import_service.py: per-file E501 ignore (data file with long strings) All 439 unit tests pass. ruff check app/ → All checks passed! Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
69 lines
1.7 KiB
Python
69 lines
1.7 KiB
Python
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import declarative_base, sessionmaker
|
|
|
|
Base = declarative_base()
|
|
|
|
# Engine and session factory are created lazily so that tests can
|
|
# override DATABASE_URL via environment *before* any import triggers
|
|
# the real PostgreSQL engine creation (which requires psycopg2).
|
|
_engine = None
|
|
_SessionLocal = None
|
|
|
|
|
|
def _get_engine():
|
|
global _engine
|
|
if _engine is None:
|
|
from app.config import settings
|
|
|
|
url = settings.DATABASE_URL
|
|
kwargs: dict = {}
|
|
if url.startswith("postgresql"):
|
|
kwargs.update(
|
|
pool_size=20,
|
|
max_overflow=10,
|
|
pool_recycle=3600,
|
|
pool_pre_ping=True,
|
|
)
|
|
_engine = create_engine(url, **kwargs)
|
|
return _engine
|
|
|
|
|
|
def _get_session_factory():
|
|
global _SessionLocal
|
|
if _SessionLocal is None:
|
|
_SessionLocal = sessionmaker(
|
|
autocommit=False, autoflush=False, bind=_get_engine()
|
|
)
|
|
return _SessionLocal
|
|
|
|
|
|
class _LazySessionLocal:
|
|
"""Proxy so ``SessionLocal()`` keeps working as before but the real
|
|
sessionmaker is only created on first call."""
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
return _get_session_factory()(*args, **kwargs)
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(_get_session_factory(), name)
|
|
|
|
|
|
SessionLocal = _LazySessionLocal()
|
|
|
|
|
|
class _EngineProxy:
|
|
"""Thin proxy so ``from app.database import engine`` still works."""
|
|
def __getattr__(self, name):
|
|
return getattr(_get_engine(), name)
|
|
|
|
|
|
engine = _EngineProxy() # type: ignore[assignment]
|
|
|
|
|
|
def get_db():
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|