mmisp-lib 0.1.5__tar.gz → 0.1.7__tar.gz

{mmisp_lib-0.1.5 → mmisp_lib-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mmisp-lib
-Version: 0.1.5
+Version: 0.1.7
 Requires-Python: >=3.11.0
 Description-Content-Type: text/markdown
 License-File: LICENSE

{mmisp_lib-0.1.5 → mmisp_lib-0.1.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mmisp-lib"
-version = "0.1.5"
+version = "0.1.7"
 description = ""
 authors = []
 readme = "README.md"

mmisp_lib-0.1.7/src/mmisp/db/all_models.py

@@ -0,0 +1,15 @@
+import importlib
+import importlib.resources
+import itertools
+
+model_pkg = "mmisp.db.models"
+all_models = (
+    resource.name[:-3]
+    for resource in importlib.resources.files(model_pkg).iterdir()
+    if resource.is_file() and resource.name != "__init__.py"
+)
+
+model_module_names = map(".".join, zip(itertools.repeat(model_pkg), all_models))
+
+for m in model_module_names:
+    importlib.import_module(m)

{mmisp_lib-0.1.5 → mmisp_lib-0.1.7}/src/mmisp/db/config.py

@@ -7,9 +7,10 @@ from dotenv import load_dotenv
 @dataclass
 class DatabaseConfig:
     DATABASE_URL: str
+    DEBUG: bool
 
 
 load_dotenv(getenv("ENV_FILE", ".env"))
 
 
-config: DatabaseConfig = DatabaseConfig(DATABASE_URL=getenv("DATABASE_URL", ""))
+config: DatabaseConfig = DatabaseConfig(DATABASE_URL=getenv("DATABASE_URL", ""), DEBUG=bool(getenv("DEBUG", False)))

mmisp_lib-0.1.7/src/mmisp/db/database.py

@@ -0,0 +1,89 @@
+import contextlib
+from typing import AsyncIterator, Self, TypeAlias
+
+from sqlalchemy.engine.url import make_url
+from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine, AsyncSession, create_async_engine
+from sqlalchemy.orm import declarative_base, sessionmaker
+
+from mmisp.db.config import config
+
+Session: TypeAlias = AsyncSession
+
+# , poolclass=NullPool)
+# async_session = sessionmaker(autoflush=False, expire_on_commit=False, class_=AsyncSession, bind=engine)
+
+Base = declarative_base()
+
+
+class DatabaseSessionManager:
+    def __init__(self: Self) -> None:
+        self._engine: AsyncEngine | None = None
+        self._sessionmaker: sessionmaker | None = None
+
+        self._url = make_url(config.DATABASE_URL)
+
+    def init(self: Self) -> None:
+        if config.DEBUG:
+            self._engine = create_async_engine(self._url, echo=True)
+        else:
+            self._engine = create_async_engine(self._url)
+        self._sessionmaker = sessionmaker(
+            autocommit=False, expire_on_commit=False, bind=self._engine, class_=AsyncSession
+        )
+
+    async def close(self: Self) -> None:
+        if self._engine is None:
+            raise Exception("DatabaseSessionManager is not initialized")
+        await self._engine.dispose()
+        self._engine = None
+        self._sessionmaker = None
+
+    @contextlib.asynccontextmanager
+    async def connect(self: Self) -> AsyncIterator[AsyncConnection]:
+        if self._engine is None:
+            raise Exception("DatabaseSessionManager is not initialized")
+
+        async with self._engine.begin() as connection:
+            try:
+                yield connection
+            except Exception:
+                await connection.rollback()
+                raise
+
+    @contextlib.asynccontextmanager
+    async def session(self: Self) -> AsyncIterator[AsyncSession]:
+        if self._sessionmaker is None:
+            raise Exception("DatabaseSessionManager is not initialized")
+
+        session = self._sessionmaker()
+        try:
+            yield session
+        except Exception:
+            await session.rollback()
+            raise
+        finally:
+            await session.close()
+
+    # Used for testing
+    async def create_all(self: Self, engine: AsyncEngine | None = None) -> None:
+        if engine is None:
+            engine = self._engine
+        async with engine.begin() as conn:
+            await conn.run_sync(Base.metadata.create_all)
+
+    async def drop_all(self: Self, engine: AsyncEngine | None = None) -> None:
+        if engine is None:
+            engine = self._engine
+        await engine.run_sync(Base.metadata.drop_all)
+
+
+async def get_db() -> Session:
+    async with sessionmanager.session() as session:
+        yield session
+
+
+async def create_all_models() -> None:
+    await sessionmanager.create_all()
+
+
+sessionmanager = DatabaseSessionManager()

{mmisp_lib-0.1.5 → mmisp_lib-0.1.7}/src/mmisp/db/models/object.py

@@ -1,10 +1,11 @@
 from sqlalchemy import Boolean, Column, ForeignKey, Integer, String
 
 from mmisp.db.database import Base
+from mmisp.db.mixins import DictMixin
 from mmisp.util.uuid import uuid
 
 
-class Object(Base):
+class Object(Base, DictMixin):
     __tablename__ = "objects"
 
     id = Column(Integer, primary_key=True, nullable=False)
@@ -24,7 +25,7 @@ class Object(Base):
     last_seen = Column(Integer, index=True, default=None)
 
 
-class ObjectTemplate(Base):
+class ObjectTemplate(Base, DictMixin):
     __tablename__ = "object_templates"
 
     id = Column(Integer, primary_key=True, nullable=False)

{mmisp_lib-0.1.5 → mmisp_lib-0.1.7}/src/mmisp_lib.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mmisp-lib
-Version: 0.1.5
+Version: 0.1.7
 Requires-Python: >=3.11.0
 Description-Content-Type: text/markdown
 License-File: LICENSE

{mmisp_lib-0.1.5 → mmisp_lib-0.1.7}/src/mmisp_lib.egg-info/SOURCES.txt

@@ -133,6 +133,7 @@ src/mmisp/api_schemas/warninglists/toggle_enable_warninglists_body.py
 src/mmisp/api_schemas/warninglists/toggle_enable_warninglists_response.py
 src/mmisp/api_schemas/warninglists/warninglist_response.py
 src/mmisp/db/__init__.py
+src/mmisp/db/all_models.py
 src/mmisp/db/config.py
 src/mmisp/db/database.py
 src/mmisp/db/mixins.py

mmisp_lib-0.1.5/src/mmisp/db/database.py

@@ -1,42 +0,0 @@
-from functools import wraps
-from typing import Any, Callable, TypeAlias
-
-from sqlalchemy.engine.url import make_url
-from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
-from sqlalchemy.orm import declarative_base, sessionmaker
-from sqlalchemy.pool import NullPool
-
-from mmisp.db.config import config
-
-Session: TypeAlias = AsyncSession
-
-url = make_url(config.DATABASE_URL)
-engine = create_async_engine(url, poolclass=NullPool)
-async_session = sessionmaker(autoflush=False, expire_on_commit=False, class_=AsyncSession, bind=engine)
-
-Base = declarative_base()
-
-
-def get_db() -> Session:
-    return async_session()
-
-
-def with_session_management(fn: Callable) -> Callable:
-    @wraps(fn)
-    async def wrapper(*args, **kwargs) -> Any:
-        db: Session = kwargs.pop("db")
-        output: Any = None
-
-        try:
-            output = await fn(*args, **kwargs, db=db)
-        finally:
-            await db.close()
-
-        return output
-
-    return wrapper
-
-
-async def create_all_models() -> None:
-    async with engine.begin() as db:
-        await db.run_sync(Base.metadata.create_all)