icij-common 0.1.1__tar.gz

icij_common-0.1.1/PKG-INFO

@@ -0,0 +1,13 @@
+Metadata-Version: 2.1
+Name: icij-common
+Version: 0.1.1
+Summary: Common utils for ICIJ libs
+Author: Clément Doumouro
+Author-email: cdoumouro@icij.org
+Requires-Python: >=3.9,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: pydantic (>=1.10,<2.0)

icij_common-0.1.1/icij_common/logging_utils.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+import contextlib
+import logging
+import sys
+from abc import ABC
+from datetime import datetime
+from functools import wraps
+from typing import List, Optional, final
+
+TRACE = 5
+logging.TRACE = TRACE
+logging.addLevelName(logging.TRACE, "TRACE")
+logging.Logger.trace = lambda inst, msg, *args, **kwargs: inst.log(
+    logging.TRACE, msg, *args, **kwargs
+)
+
+
+class DifferedLoggingMessage:
+    def __init__(self, fn, *args, **kwargs):
+        self.fn = fn
+        self.args = args
+        self.kwargs = kwargs
+
+    def __str__(self):
+        return str(self.fn(*self.args, **self.kwargs))
+
+
+class LogWithNameMixin(ABC):
+    def __init__(self, logger: logging.Logger):
+        self._logger = logger
+
+    @final
+    def info(self, msg, *args, **kwargs):
+        self._logger.info(msg, *args, **kwargs)
+
+    @final
+    def debug(self, msg, *args, **kwargs):
+        self._logger.debug(msg, *args, **kwargs)
+
+    @final
+    def error(self, msg, *args, **kwargs):
+        self._logger.error(msg, *args, **kwargs)
+
+
+class WorkerIdFilter(logging.Filter):
+    def __init__(self, worker_id: str):
+        super().__init__()
+        self._worker_id = worker_id
+
+    def filter(self, record):
+        record.workerid = self._worker_id
+        return True
+
+
+def log_elapsed_time(
+    logger: logging.Logger, level: int, output_msg: Optional[str] = None
+):
+    if output_msg is None:
+        output_msg = "Elapsed time ->:\n{elapsed_time}"
+
+    def get_wrapper(fn):
+        @wraps(fn)
+        def wrapped(*args, **kwargs):
+            start = datetime.now()
+            msg_fmt = dict()
+            res = fn(*args, **kwargs)
+            if "elapsed_time" in output_msg:
+                msg_fmt["elapsed_time"] = datetime.now() - start
+            logger.log(level, output_msg.format(**msg_fmt))
+            return res
+
+        return wrapped
+
+    return get_wrapper
+
+
+@contextlib.contextmanager
+def log_elapsed_time_cm(
+    logger: logging.Logger, level: int, output_msg: Optional[str] = None
+):
+    if output_msg is None:
+        output_msg = "Elapsed time ->:\n{elapsed_time}"
+    start = datetime.now()
+    yield
+    end = datetime.now() - start
+    msg_fmt = dict()
+    if "elapsed_time" in output_msg:
+        msg_fmt["elapsed_time"] = end
+    logger.log(level, output_msg.format(**msg_fmt))
+
+
+STREAM_HANDLER_FMT = "[%(levelname)s][%(asctime)s.%(msecs)03d][%(name)s]: %(message)s"
+STREAM_HANDLER_FMT_WITH_WORKER_ID = (
+    "[%(levelname)s][%(asctime)s.%(msecs)03d][%(workerid)s][%(name)s]: %(message)s"
+)
+DATE_FMT = "%H:%M:%S"
+
+
+def setup_loggers(loggers: List[str]):
+    level = logging.INFO
+    stream_handler = logging.StreamHandler(sys.stderr)
+    stream_handler.setFormatter(logging.Formatter(STREAM_HANDLER_FMT, DATE_FMT))
+    for logger in loggers:
+        logger = logging.getLogger(logger)
+        logger.setLevel(level)
+        logger.handlers = []
+        logger.addHandler(stream_handler)

icij_common-0.1.1/icij_common/neo4j/constants.py

@@ -0,0 +1,38 @@
+PROJECT_REGISTRY_DB = "datashare-project-registry"
+
+PROJECT_RUNS_MIGRATION = "_RUNS"
+PROJECT_NAME = "name"
+PROJECT_NODE = "_Project"
+
+MIGRATION_COMPLETED = "completed"
+MIGRATION_LABEL = "label"
+MIGRATION_NODE = "_Migration"
+MIGRATION_PROJECT = "project"
+MIGRATION_STARTED = "started"
+MIGRATION_STATUS = "status"
+MIGRATION_VERSION = "version"
+
+TASK_NODE = "_Task"
+TASK_COMPLETED_AT = "completedAt"
+TASK_CREATED_AT = "createdAt"
+TASK_ID = "id"
+TASK_INPUTS = "inputs"
+TASK_PROGRESS = "progress"
+TASK_RETRIES = "retries"
+TASK_TYPE = "type"
+
+TASK_LOCK_NODE = "_TaskLock"
+TASK_LOCK_TASK_ID = "taskId"
+TASK_LOCK_WORKER_ID = "workerId"
+
+TASK_ERROR_NODE = "_TaskError"
+TASK_ERROR_DETAIL = "detail"
+TASK_ERROR_ID = "id"
+TASK_ERROR_OCCURRED_AT = "occurredAt"
+TASK_ERROR_TITLE = "title"
+
+TASK_ERROR_OCCURRED_TYPE = "_OCCURRED_DURING"
+
+TASK_RESULT_NODE = "_TaskResult"
+TASK_HAS_RESULT_TYPE = "_HAS_RESULT"
+TASK_RESULT_RESULT = "result"

icij_common-0.1.1/icij_common/neo4j/migrate.py

@@ -0,0 +1,326 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+from collections.abc import Coroutine
+from datetime import datetime
+from distutils.version import StrictVersion
+from enum import Enum, unique
+from inspect import signature
+from typing import Any, Callable, List, Optional, Sequence, Union
+
+import neo4j
+from neo4j.exceptions import ConstraintError
+
+from icij_common.neo4j.constants import (
+    MIGRATION_COMPLETED,
+    MIGRATION_LABEL,
+    MIGRATION_NODE,
+    MIGRATION_PROJECT,
+    MIGRATION_STARTED,
+    MIGRATION_STATUS,
+    MIGRATION_VERSION,
+)
+from icij_common.pydantic_utils import NoEnumModel
+from .projects import (
+    Project,
+    create_project_db,
+    create_project_tx,
+    project_db_session,
+    projects_tx,
+    registry_db_session,
+)
+
+logger = logging.getLogger(__name__)
+
+TransactionFn = Callable[[neo4j.AsyncTransaction], Coroutine]
+ExplicitTransactionFn = Callable[[neo4j.Session], Coroutine]
+MigrationFn = Union[TransactionFn, ExplicitTransactionFn]
+
+_MIGRATION_TIMEOUT_MSG = """Migration timeout expired !
+Please check that a migration is indeed in progress. If the application is in a \
+deadlock restart it forcing the migration index cleanup."""
+
+
+class MigrationError(RuntimeError):
+    pass
+
+
+@unique
+class MigrationStatus(str, Enum):
+    IN_PROGRESS = "IN_PROGRESS"
+    DONE = "DONE"
+
+
+class MigrationVersion(StrictVersion):
+    @classmethod
+    def __get_validators__(cls):
+        def validator(v: Any) -> MigrationVersion:
+            if isinstance(v, (str, MigrationVersion)):
+                return MigrationVersion(v)
+            raise ValueError(
+                f"Must be a {MigrationVersion.__name__} or a {str.__name__}, "
+                f"found {type(v)}"
+            )
+
+        yield validator
+
+
+class _BaseMigration(NoEnumModel):
+    version: MigrationVersion
+    label: str
+    status: MigrationStatus = MigrationStatus.IN_PROGRESS
+
+
+class Neo4jMigration(_BaseMigration):
+    # It would have been cleaner to create
+    # (p:_Project)-[:_RUNS { id: p.name + m.version }]->(m:_Migration)
+    # relationships. However, neo4j < 5.7 doesn't support unique constraint on
+    # relationships properties which prevents from implementing the locking mechanism
+    # properly. We hence enforce unique constraint on
+    # (_Migration.version, _Migration.project)
+    project: str
+    started: datetime
+    completed: Optional[datetime] = None
+    status: MigrationStatus = MigrationStatus.IN_PROGRESS
+
+    @classmethod
+    def from_neo4j(cls, record: neo4j.Record, key="migration") -> Neo4jMigration:
+        migration = dict(record.value(key))
+        if "started" in migration:
+            migration["started"] = migration["started"].to_native()
+        if "completed" in migration:
+            migration["completed"] = migration["completed"].to_native()
+        return Neo4jMigration(**migration)
+
+
+class Migration(_BaseMigration):
+    migration_fn: MigrationFn
+
+
+MigrationRegistry: Sequence[Migration]
+
+
+async def _migrate_with_lock(
+    *,
+    project_session: neo4j.AsyncSession,
+    registry_session: neo4j.AsyncSession,
+    project: str,
+    migration: Migration,
+):
+    # Note: all migrations.py should be carefully tested otherwise they will lock
+    # the DB...
+
+    # Lock the DB first, raising in case a migration already exists
+    logger.debug("Trying to run migration to %s...", migration.label)
+    await registry_session.execute_write(
+        create_migration_tx,
+        project=project,
+        migration_version=str(migration.version),
+        migration_label=migration.label,
+    )
+    # Then run to migration
+    logger.debug("Acquired write lock to %s !", migration.label)
+    sig = signature(migration.migration_fn)
+    first_param = list(sig.parameters)[0]
+    if first_param == "tx":
+        await project_session.execute_write(migration.migration_fn)
+    elif first_param == "sess":
+        await migration.migration_fn(project_session)
+    else:
+        raise ValueError(f"Invalid migration function: {migration.migration_fn}")
+    # Finally free the lock
+    await registry_session.execute_write(
+        complete_migration_tx,
+        project=project,
+        migration_version=str(migration.version),
+    )
+    logger.debug("Marked %s as complete !", migration.label)
+
+
+async def create_migration_tx(
+    tx: neo4j.AsyncTransaction,
+    *,
+    project: str,
+    migration_version: str,
+    migration_label: str,
+) -> Neo4jMigration:
+    query = f"""CREATE (m:{MIGRATION_NODE} {{
+    {MIGRATION_PROJECT}: $project,
+    {MIGRATION_LABEL}: $label,
+    {MIGRATION_VERSION}: $version,
+    {MIGRATION_STATUS}: $status,
+    {MIGRATION_STARTED}: $started
+}})
+RETURN m as migration"""
+    res = await tx.run(
+        query,
+        label=migration_label,
+        version=migration_version,
+        project=project,
+        status=MigrationStatus.IN_PROGRESS.value,
+        started=datetime.now(),
+    )
+    migration = await res.single()
+    if migration is None:
+        raise ValueError(f"Couldn't find migration {migration_version} for {project}")
+    migration = Neo4jMigration.from_neo4j(migration)
+    return migration
+
+
+async def complete_migration_tx(
+    tx: neo4j.AsyncTransaction, *, project: str, migration_version: str
+) -> Neo4jMigration:
+    query = f"""MATCH (m:{MIGRATION_NODE} {{
+        {MIGRATION_VERSION}: $version,
+        {MIGRATION_PROJECT}: $project
+     }})
+SET m += {{ {MIGRATION_STATUS}: $status, {MIGRATION_COMPLETED}: $completed }} 
+RETURN m as migration"""
+    res = await tx.run(
+        query,
+        version=migration_version,
+        project=project,
+        status=MigrationStatus.DONE.value,
+        completed=datetime.now(),
+    )
+    migration = await res.single()
+    migration = Neo4jMigration.from_neo4j(migration)
+    return migration
+
+
+async def project_migrations_tx(
+    tx: neo4j.AsyncTransaction, project: str
+) -> List[Neo4jMigration]:
+    query = f"""MATCH (m:{MIGRATION_NODE} {{ {MIGRATION_PROJECT}: $project }})
+RETURN m as migration
+"""
+    res = await tx.run(query, project=project)
+    migrations = [Neo4jMigration.from_neo4j(rec) async for rec in res]
+    return migrations
+
+
+async def delete_all_migrations(driver: neo4j.AsyncDriver):
+    query = f"""MATCH (m:{MIGRATION_NODE})
+DETACH DELETE m"""
+    async with registry_db_session(driver) as sess:
+        await sess.run(query)
+
+
+async def retrieve_projects(neo4j_driver: neo4j.AsyncDriver) -> List[Project]:
+    async with registry_db_session(neo4j_driver) as sess:
+        projects = await sess.execute_read(projects_tx)
+    return projects
+
+
+async def migrate_db_schemas(
+    neo4j_driver: neo4j.AsyncDriver,
+    registry: MigrationRegistry,
+    *,
+    timeout_s: float,
+    throttle_s: float,
+):
+    projects = await retrieve_projects(neo4j_driver)
+    tasks = [
+        migrate_project_db_schema(
+            neo4j_driver,
+            registry,
+            project=p.name,
+            timeout_s=timeout_s,
+            throttle_s=throttle_s,
+        )
+        for p in projects
+    ]
+    await asyncio.gather(*tasks)
+
+
+async def migrate_project_db_schema(
+    neo4j_driver: neo4j.AsyncDriver,
+    registry: MigrationRegistry,
+    project: str,
+    *,
+    timeout_s: float,
+    throttle_s: float,
+):
+    logger.info("Migrating project %s", project)
+    start = time.monotonic()
+    if not registry:
+        return
+    todo = sorted(registry, key=lambda m: m.version)
+    async with registry_db_session(neo4j_driver) as registry_sess:
+        async with project_db_session(neo4j_driver, project=project) as project_sess:
+            while "Waiting for DB to be migrated or for a timeout":
+                migrations = await registry_sess.execute_read(
+                    project_migrations_tx, project=project
+                )
+                in_progress = [
+                    m for m in migrations if m.status is MigrationStatus.IN_PROGRESS
+                ]
+                if len(in_progress) > 1:
+                    raise MigrationError(
+                        f"Found several migration in progress: {in_progress}"
+                    )
+                if in_progress:
+                    logger.info(
+                        "Found that %s is in progress, waiting for %s seconds...",
+                        in_progress[0].label,
+                        throttle_s,
+                    )
+                    await asyncio.sleep(throttle_s)
+                else:
+                    done = [m for m in migrations if m.status is MigrationStatus.DONE]
+                    if done:
+                        current_version = max((m.version for m in done))
+                        todo = [m for m in todo if m.version > current_version]
+                    if not todo:
+                        break
+                    try:
+                        await _migrate_with_lock(
+                            project_session=project_sess,
+                            registry_session=registry_sess,
+                            project=project,
+                            migration=todo[0],
+                        )
+                        todo = todo[1:]
+                        continue
+                    except ConstraintError:
+                        logger.info(
+                            "Migration %s has just started somewhere else, "
+                            " waiting for %s seconds...",
+                            todo[0].label,
+                            throttle_s,
+                        )
+                        await asyncio.sleep(throttle_s)
+                elapsed = time.monotonic() - start
+                if elapsed > timeout_s:
+                    logger.error(_MIGRATION_TIMEOUT_MSG)
+                    raise MigrationError(_MIGRATION_TIMEOUT_MSG)
+                continue
+
+
+async def init_project(
+    neo4j_driver: neo4j.AsyncDriver,
+    name: str,
+    registry: MigrationRegistry,
+    *,
+    timeout_s: float,
+    throttle_s: float,
+) -> bool:
+    # Create project DB
+    await create_project_db(neo4j_driver, project=name)
+
+    # Create project
+    async with registry_db_session(neo4j_driver) as sess:
+        project, already_exists = await sess.execute_write(create_project_tx, name=name)
+
+    # Migrate it
+    await migrate_project_db_schema(
+        neo4j_driver,
+        registry=registry,
+        project=project.name,
+        timeout_s=timeout_s,
+        throttle_s=throttle_s,
+    )
+
+    return already_exists

icij_common-0.1.1/icij_common/neo4j/projects.py

@@ -0,0 +1,168 @@
+from __future__ import annotations
+import logging
+from contextlib import asynccontextmanager
+from distutils.version import StrictVersion
+from typing import AsyncGenerator, List, Optional, Tuple
+
+import neo4j
+
+from icij_common.neo4j.constants import (
+    MIGRATION_NODE,
+    MIGRATION_PROJECT,
+    MIGRATION_VERSION,
+    PROJECT_NODE,
+    PROJECT_REGISTRY_DB,
+)
+
+from icij_common.neo4j.constants import PROJECT_NAME
+
+from icij_common.pydantic_utils import ICIJModel
+
+logger = logging.getLogger(__name__)
+
+NEO4J_COMMUNITY_DB = "neo4j"
+_IS_ENTERPRISE: Optional[bool] = None
+_NEO4J_VERSION: Optional[StrictVersion] = None
+_SUPPORTS_PARALLEL: Optional[bool] = None
+
+_COMPONENTS_QUERY = """CALL dbms.components() YIELD versions, edition
+RETURN versions, edition"""
+
+
+class Project(ICIJModel):
+    name: str
+
+    @classmethod
+    def from_neo4j(cls, record: neo4j.Record, key="project") -> Project:
+        project = dict(record.value(key))
+        return Project(**project)
+
+
+async def create_project_registry_db(neo4j_driver: neo4j.AsyncDriver):
+    if await is_enterprise(neo4j_driver):
+        logger.info("Creating project registry DB...")
+        query = "CREATE DATABASE $registry_db IF NOT EXISTS"
+        await neo4j_driver.execute_query(query, registry_db=PROJECT_REGISTRY_DB)
+    else:
+        logger.info("Using default db as registry DB !")
+
+
+async def projects_tx(tx: neo4j.AsyncTransaction) -> List[Project]:
+    query = f"MATCH (project:{PROJECT_NODE}) RETURN project"
+    res = await tx.run(query)
+    projects = [Project.from_neo4j(p) async for p in res]
+    return projects
+
+
+async def create_project_db(neo4j_driver: neo4j.AsyncDriver, project: str):
+    if await is_enterprise(neo4j_driver):
+        db_name = await project_db(neo4j_driver, project=project)
+        query = "CREATE DATABASE $db_name IF NOT EXISTS"
+        await neo4j_driver.execute_query(query, db_name=db_name)
+
+
+async def add_project_support_migration_tx(tx: neo4j.AsyncTransaction):
+    await create_project_unique_name_constraint_tx(tx)
+    await create_migration_unique_project_and_version_constraint_tx(tx)
+
+
+async def create_project_unique_name_constraint_tx(tx: neo4j.AsyncTransaction):
+    constraint_query = f"""CREATE CONSTRAINT constraint_project_unique_name
+IF NOT EXISTS
+FOR (p:{PROJECT_NODE})
+REQUIRE (p.{PROJECT_NAME}) IS UNIQUE
+"""
+    await tx.run(constraint_query)
+
+
+async def create_migration_unique_project_and_version_constraint_tx(
+    tx: neo4j.AsyncTransaction,
+):
+    constraint_query = f"""CREATE CONSTRAINT
+     constraint_migration_unique_project_and_version
+IF NOT EXISTS 
+FOR (m:{MIGRATION_NODE})
+REQUIRE (m.{MIGRATION_VERSION}, m.{MIGRATION_PROJECT}) IS UNIQUE
+"""
+    await tx.run(constraint_query)
+
+
+async def create_project_tx(
+    tx: neo4j.AsyncTransaction, name: str
+) -> Tuple[Project, bool]:
+    if name == PROJECT_REGISTRY_DB:
+        raise ValueError(
+            f'Bad luck, name "{PROJECT_REGISTRY_DB}" is reserved for internal use.'
+            f" Can't initialize project"
+        )
+    query = f"""MERGE (project:{PROJECT_NODE} {{ {PROJECT_NAME}: $name }})
+RETURN project"""
+    res = await tx.run(query, name=name)
+    rec = await res.single()
+    summary = await res.consume()
+    existed = summary.counters.nodes_created == 0
+    project = Project.from_neo4j(rec)
+    return project, existed
+
+
+async def project_registry_db(neo4j_driver: neo4j.AsyncDriver) -> str:
+    if await is_enterprise(neo4j_driver):
+        return PROJECT_REGISTRY_DB
+    return NEO4J_COMMUNITY_DB
+
+
+async def project_db(neo4j_driver: neo4j.AsyncDriver, project: str) -> str:
+    if await is_enterprise(neo4j_driver):
+        return project
+    return NEO4J_COMMUNITY_DB
+
+
+def project_index(project: str) -> str:
+    return project
+
+
+@asynccontextmanager
+async def project_db_session(
+    neo4j_driver: neo4j.AsyncDriver, project: str
+) -> AsyncGenerator[neo4j.AsyncSession, None]:
+    db = await project_db(neo4j_driver, project)
+    sess_ctx = neo4j_driver.session(database=db)
+    async with sess_ctx as sess:
+        yield sess
+
+
+@asynccontextmanager
+async def registry_db_session(neo4j_driver: neo4j.AsyncDriver) -> neo4j.AsyncSession:
+    session = neo4j_driver.session(database=await project_registry_db(neo4j_driver))
+    async with session as sess:
+        yield sess
+
+
+async def _get_components(neo4j_driver: neo4j.AsyncDriver):
+    async with neo4j_driver.session(database=neo4j.SYSTEM_DATABASE) as sess:
+        res = await sess.run(_COMPONENTS_QUERY)
+        res = await res.single()
+    global _IS_ENTERPRISE
+    global _NEO4J_VERSION
+    _IS_ENTERPRISE = res["edition"] != "community"
+    _NEO4J_VERSION = StrictVersion(res["versions"][0])
+
+
+async def server_version(neo4j_driver: neo4j.AsyncDriver) -> StrictVersion:
+    if _NEO4J_VERSION is None:
+        await _get_components(neo4j_driver)
+    return _NEO4J_VERSION
+
+
+async def supports_parallel_runtime(neo4j_driver: neo4j.AsyncDriver) -> bool:
+    global _SUPPORTS_PARALLEL
+    if _SUPPORTS_PARALLEL is None:
+        version = await server_version(neo4j_driver)
+        _SUPPORTS_PARALLEL = version >= StrictVersion("5.13")
+    return _SUPPORTS_PARALLEL
+
+
+async def is_enterprise(neo4j_driver: neo4j.AsyncDriver) -> bool:
+    if _IS_ENTERPRISE is None:
+        await _get_components(neo4j_driver)
+    return _IS_ENTERPRISE

icij_common-0.1.1/icij_common/neo4j/test_utils.py

@@ -0,0 +1,129 @@
+# pylint: disable=redefined-outer-name
+import contextlib
+from typing import AsyncGenerator
+
+import neo4j
+import pytest
+from neo4j import AsyncGraphDatabase
+
+import icij_common
+from icij_common.neo4j.projects import NEO4J_COMMUNITY_DB
+
+NEO4J_TEST_PORT = 7688
+NEO4J_TEST_USER = "neo4j"
+NEO4J_TEST_PASSWORD = "theneo4jpassword"
+NEO4J_TEST_AUTH = neo4j.basic_auth(NEO4J_TEST_USER, NEO4J_TEST_PASSWORD)
+
+
+async def wipe_db(session: neo4j.AsyncSession):
+    # Indices and constraints
+    query = "CALL apoc.schema.assert({}, {})"
+    await session.run(query)
+    # Documents
+    query = """MATCH (n)
+DETACH DELETE n
+    """
+    await session.run(query)
+
+
+@contextlib.asynccontextmanager
+async def _build_neo4j_driver():
+    uri = f"neo4j://127.0.0.1:{NEO4J_TEST_PORT}"
+    async with AsyncGraphDatabase.driver(  # pylint: disable=not-async-context-manager
+        uri, auth=NEO4J_TEST_AUTH
+    ) as driver:
+        yield driver
+
+
+@pytest.fixture(scope="module")
+async def neo4j_test_driver_module() -> AsyncGenerator[neo4j.AsyncDriver, None]:
+    async with _build_neo4j_driver() as driver:
+        async with driver.session(database=neo4j.DEFAULT_DATABASE) as sess:
+            await wipe_db(sess)
+        yield driver
+
+
+@pytest.fixture(scope="session")
+async def neo4j_test_driver_session() -> AsyncGenerator[neo4j.AsyncDriver, None]:
+    async with _build_neo4j_driver() as driver:
+        async with driver.session(database=neo4j.DEFAULT_DATABASE) as sess:
+            await wipe_db(sess)
+        yield driver
+
+
+@pytest.fixture()
+async def neo4j_test_driver() -> AsyncGenerator[neo4j.AsyncDriver, None]:
+    async with _build_neo4j_driver() as driver:
+        async with driver.session(database="neo4j") as sess:
+            await wipe_db(sess)
+        yield driver
+
+
+@pytest.fixture(scope="session")
+async def neo4j_test_session_session(
+    neo4j_test_driver_session: neo4j.Driver,
+) -> AsyncGenerator[neo4j.AsyncSession, None]:
+    driver = neo4j_test_driver_session
+    async with driver.session(database="neo4j") as sess:
+        await wipe_db(sess)
+        yield sess
+
+
+@pytest.fixture(scope="module")
+async def neo4j_test_session_module(
+    neo4j_test_session_session: neo4j.AsyncSession,
+) -> neo4j.AsyncSession:
+    session = neo4j_test_session_session
+    await wipe_db(session)
+    return session
+
+
+@pytest.fixture()
+async def neo4j_test_session(
+    neo4j_test_session_session: neo4j.AsyncSession,
+) -> neo4j.AsyncSession:
+    session = neo4j_test_session_session
+    await wipe_db(session)
+    return session
+
+
+async def mocked_is_enterprise(_: neo4j.AsyncDriver) -> bool:
+    return True
+
+
+@contextlib.asynccontextmanager
+async def _mocked_project_db_session(
+    neo4j_driver: neo4j.AsyncDriver, project: str  # pylint: disable=unused-argument
+) -> neo4j.AsyncSession:
+    async with neo4j_driver.session(database=NEO4J_COMMUNITY_DB) as sess:
+        yield sess
+
+
+async def _mocked_project_registry_db(
+    neo4j_driver: neo4j.AsyncDriver,  # pylint: disable=unused-argument
+) -> str:
+    return NEO4J_COMMUNITY_DB
+
+
+@pytest.fixture()
+def mock_enterprise(monkeypatch):
+    mock_enterprise_(monkeypatch)
+
+
+def mock_enterprise_(monkeypatch):
+    monkeypatch.setattr(
+        icij_common.neo4j.projects, "project_db_session", _mocked_project_db_session
+    )
+    monkeypatch.setattr(
+        icij_common.neo4j.migrate,
+        "project_db_session",
+        _mocked_project_db_session,
+    )
+    monkeypatch.setattr(
+        icij_common.neo4j.projects,
+        "project_registry_db",
+        _mocked_project_registry_db,
+    )
+    monkeypatch.setattr(
+        icij_common.neo4j.projects, "is_enterprise", mocked_is_enterprise
+    )

icij_common-0.1.1/icij_common/pydantic_utils.py

@@ -0,0 +1,239 @@
+import dataclasses
+from collections import defaultdict, deque
+from datetime import datetime
+from enum import Enum
+from functools import cached_property
+from pathlib import PurePath
+from types import GeneratorType
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union, cast
+
+from pydantic import BaseModel
+from pydantic.fields import FieldInfo
+from pydantic.json import ENCODERS_BY_TYPE
+
+SetIntStr = Set[Union[int, str]]
+DictIntStrAny = Dict[Union[int, str], Any]
+
+
+def to_lower_camel(field: str) -> str:
+    return "".join(
+        w.capitalize() if i > 0 else w for i, w in enumerate(field.split("_"))
+    )
+
+
+_FIELD_ARGS = ["include", "exclude", "update"]
+
+_SCHEMAS = dict()
+
+
+# TODO: remove this one when migrating to pydantic 2.0
+def safe_copy(model: BaseModel, **kwargs):
+    if model.__class__ not in _SCHEMAS:
+        _SCHEMAS[model.__class__] = dict()
+        # Model.copy is always without alias
+        _SCHEMAS[model.__class__] = model.__class__.schema(by_alias=False)
+    schema = _SCHEMAS[model.__class__]
+    for k in _FIELD_ARGS:
+        if k not in kwargs:
+            continue
+        for field in kwargs[k]:
+            if field not in schema["properties"]:
+                msg = f'Unknown attribute "{field}" for {model.__class__}'
+                raise AttributeError(msg)
+
+    return cast(model.__class__, model.copy(**kwargs))
+
+
+class ICIJModel(BaseModel):
+    class Config:
+        allow_mutation = False
+        extra = "forbid"
+        allow_population_by_field_name = True
+        keep_untouched = (cached_property,)
+        use_enum_values = True
+
+
+def get_field_default_value(attr: FieldInfo):
+    # Ugly work around pydantic v1 limitations on Field and default values
+    if isinstance(attr, FieldInfo):
+        if attr.default_factory is not None:
+            return attr.default_factory()
+        return attr.default
+    return attr
+
+
+class LowerCamelCaseModel(ICIJModel):
+    class Config:
+        alias_generator = to_lower_camel
+
+
+class IgnoreExtraModel(ICIJModel):
+    class Config:
+        extra = "ignore"
+
+
+class NoEnumModel(ICIJModel):
+    class Config:
+        use_enum_values = False
+
+
+class ISODatetime(ICIJModel):
+    class Config:
+        json_encoders = {datetime: lambda x: x.isoformat()}
+
+
+def generate_encoders_by_class_tuples(
+    type_encoder_map: Dict[Any, Callable[[Any], Any]]
+) -> Dict[Callable[[Any], Any], Tuple[Any, ...]]:
+    encoders_by_class_tuples: Dict[Callable[[Any], Any], Tuple[Any, ...]] = defaultdict(
+        tuple
+    )
+    for type_, encoder in type_encoder_map.items():
+        encoders_by_class_tuples[encoder] += (type_,)
+    return encoders_by_class_tuples
+
+
+encoders_by_class_tuples_ = generate_encoders_by_class_tuples(ENCODERS_BY_TYPE)
+
+
+def jsonable_encoder(
+    obj: Any,
+    include: Optional[Union[SetIntStr, DictIntStrAny]] = None,
+    exclude: Optional[Union[SetIntStr, DictIntStrAny]] = None,
+    by_alias: bool = True,
+    exclude_unset: bool = False,
+    exclude_defaults: bool = False,
+    exclude_none: bool = False,
+    custom_encoder: Optional[Dict[Any, Callable[[Any], Any]]] = None,
+    sqlalchemy_safe: bool = True,
+) -> Any:
+    custom_encoder = custom_encoder or {}
+    if custom_encoder:
+        if type(obj) in custom_encoder:
+            return custom_encoder[type(obj)](obj)
+        for encoder_type, encoder_instance in custom_encoder.items():
+            if isinstance(obj, encoder_type):
+                return encoder_instance(obj)
+    if include is not None and not isinstance(include, (set, dict)):
+        include = set(include)
+    if exclude is not None and not isinstance(exclude, (set, dict)):
+        exclude = set(exclude)
+    if isinstance(obj, BaseModel):
+        encoder = getattr(obj.__config__, "json_encoders", {})
+        if custom_encoder:
+            encoder.update(custom_encoder)
+        obj_dict = obj.dict(
+            include=include,
+            exclude=exclude,
+            by_alias=by_alias,
+            exclude_unset=exclude_unset,
+            exclude_none=exclude_none,
+            exclude_defaults=exclude_defaults,
+        )
+        if "__root__" in obj_dict:
+            obj_dict = obj_dict["__root__"]
+        return jsonable_encoder(
+            obj_dict,
+            exclude_none=exclude_none,
+            exclude_defaults=exclude_defaults,
+            custom_encoder=encoder,
+            sqlalchemy_safe=sqlalchemy_safe,
+        )
+    if dataclasses.is_dataclass(obj):
+        obj_dict = dataclasses.asdict(obj)
+        return jsonable_encoder(
+            obj_dict,
+            include=include,
+            exclude=exclude,
+            by_alias=by_alias,
+            exclude_unset=exclude_unset,
+            exclude_defaults=exclude_defaults,
+            exclude_none=exclude_none,
+            custom_encoder=custom_encoder,
+            sqlalchemy_safe=sqlalchemy_safe,
+        )
+    if isinstance(obj, Enum):
+        return obj.value
+    if isinstance(obj, PurePath):
+        return str(obj)
+    if isinstance(obj, (str, int, float, type(None))):
+        return obj
+    if isinstance(obj, dict):
+        encoded_dict = {}
+        allowed_keys = set(obj.keys())
+        if include is not None:
+            allowed_keys &= set(include)
+        if exclude is not None:
+            allowed_keys -= set(exclude)
+        for key, value in obj.items():
+            if (
+                (
+                    not sqlalchemy_safe
+                    or (not isinstance(key, str))
+                    or (not key.startswith("_sa"))
+                )
+                and (value is not None or not exclude_none)
+                and key in allowed_keys
+            ):
+                encoded_key = jsonable_encoder(
+                    key,
+                    by_alias=by_alias,
+                    exclude_unset=exclude_unset,
+                    exclude_none=exclude_none,
+                    custom_encoder=custom_encoder,
+                    sqlalchemy_safe=sqlalchemy_safe,
+                )
+                encoded_value = jsonable_encoder(
+                    value,
+                    by_alias=by_alias,
+                    exclude_unset=exclude_unset,
+                    exclude_none=exclude_none,
+                    custom_encoder=custom_encoder,
+                    sqlalchemy_safe=sqlalchemy_safe,
+                )
+                encoded_dict[encoded_key] = encoded_value
+        return encoded_dict
+    if isinstance(obj, (list, set, frozenset, GeneratorType, tuple, deque)):
+        encoded_list = []
+        for item in obj:
+            encoded_list.append(
+                jsonable_encoder(
+                    item,
+                    include=include,
+                    exclude=exclude,
+                    by_alias=by_alias,
+                    exclude_unset=exclude_unset,
+                    exclude_defaults=exclude_defaults,
+                    exclude_none=exclude_none,
+                    custom_encoder=custom_encoder,
+                    sqlalchemy_safe=sqlalchemy_safe,
+                )
+            )
+        return encoded_list
+
+    if type(obj) in ENCODERS_BY_TYPE:
+        return ENCODERS_BY_TYPE[type(obj)](obj)
+    for encoder, classes_tuple in encoders_by_class_tuples_.items():
+        if isinstance(obj, classes_tuple):
+            return encoder(obj)
+
+    try:
+        data = dict(obj)
+    except Exception as e:  # pylint: disable=broad-exception-caught
+        errors: List[Exception] = [e]
+        try:
+            data = vars(obj)
+        except Exception as other_e:
+            errors.append(other_e)
+            raise ValueError(errors) from other_e
+    return jsonable_encoder(
+        data,
+        include=include,
+        exclude=exclude,
+        by_alias=by_alias,
+        exclude_unset=exclude_unset,
+        exclude_defaults=exclude_defaults,
+        exclude_none=exclude_none,
+        custom_encoder=custom_encoder,
+        sqlalchemy_safe=sqlalchemy_safe,
+    )

icij_common-0.1.1/icij_common/test_utils.py

@@ -0,0 +1,75 @@
+import asyncio
+import contextlib
+import os
+import traceback
+from copy import copy
+from time import monotonic, sleep
+from typing import Awaitable, Callable, Optional
+
+import pytest
+
+TEST_PROJECT = "test_project"
+
+
+def true_after(
+    state_statement: Callable, *, after_s: float, sleep_s: float = 0.01
+) -> bool:
+    start = monotonic()
+    while "waiting for the statement to be True":
+        try:
+            assert state_statement()
+            return True
+        except AssertionError:
+            if monotonic() - start < after_s:
+                sleep(sleep_s)
+                continue
+            return False
+
+
+async def async_true_after(
+    state_statement: Callable[[], Awaitable[bool]],
+    *,
+    after_s: float,
+    sleep_s: float = 0.01,
+) -> bool:
+    start = monotonic()
+    while "waiting for the statement to be True":
+        try:
+            assert await state_statement()
+            return True
+        except AssertionError:
+            if monotonic() - start < after_s:
+                await asyncio.sleep(sleep_s)
+                continue
+            return False
+
+
+@contextlib.contextmanager
+def fail_if_exception(msg: Optional[str] = None):
+    try:
+        yield
+    except Exception as e:  # pylint: disable=W0703
+        trace = "".join(traceback.format_exception(None, e, e.__traceback__))
+        if msg is None:
+            msg = "Test failed due to the following error"
+        pytest.fail(f"{msg}\n{trace}")
+
+
+@pytest.fixture()
+def reset_env():
+    old_env = copy(dict(os.environ))
+    try:
+        yield
+    finally:
+        os.environ.clear()
+        os.environ.update(old_env)
+
+
+# Define a session level even_loop fixture to overcome limitation explained here:
+# https://github.com/tortoise/tortoise-orm/issues/638#issuecomment-830124562
+@pytest.fixture(scope="session")
+def event_loop():
+    policy = asyncio.get_event_loop_policy()
+    loop = policy.new_event_loop()
+    yield loop
+    loop.close()

icij_common-0.1.1/pyproject.toml

@@ -0,0 +1,37 @@
+[tool.poetry]
+name = "icij-common"
+version = "0.1.1"
+description = "Common utils for ICIJ libs"
+authors = [
+    "Clément Doumouro <cdoumouro@icij.org>",
+    "ICIJ <engineering@icij.org>",
+]
+# TODO: add readme
+packages = [
+    { include = "icij_common" }
+]
+exclude = ["icij_common/tests"]
+
+[tool.black]
+line-length = 88
+target-version = ["py39"]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+
+[tool.poetry.dependencies]
+python = "^3.9"
+pydantic = "^1.10" # Sadly needed for now because of Pyinstaller
+
+[tool.poetry.group.dev.dependencies]
+neo4j = "^5.0.0"
+pylint = "^3.1.0"
+pytest = "^7.2.1"
+pytest-asyncio = "^0.20.3"
+
+[tool.poetry.group.neo4j.dependencies]
+neo4j = "^5.0.0"
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"