Ops 2662 can migration 3

backend/data_tools/environment/pytest_data_tools.py

@@ -0,0 +1,33 @@
+from data_tools.environment.types import DataToolsConfig
+
+
+class PytestDataToolsConfig(DataToolsConfig):
+    @property
+    def db_connection_string(self) -> str:
+        return (
+            "postgresql://postgres:local_password@localhost:54321/postgres"  # pragma: allowlist secret
+        )
+
+    @property
+    def verbosity(self) -> bool:
+        return True
+
+    @property
+    def is_remote(self) -> bool:
+        return False
+
+    @property
+    def file_system_path(self) -> str:
+        return "."
+
+    @property
+    def vault_url(self) -> str | None:
+        return None
+
+    @property
+    def vault_file_storage_key(self) -> str | None:
+        return None
+
+    @property
+    def file_storage_auth_method(self) -> str | None:
+        return None

backend/data_tools/scripts/get_csv.sh → backend/data_tools/scripts/load_cans.sh

@@ -5,17 +5,14 @@ export PYTHONPATH=.:$PYTHONPATH
 
 ENV=$1
 INPUT_CSV=$2
-OUTPUT_CSV=$3
 
 echo "Activating virtual environment..."
 . .venv/bin/activate
 
 echo "ENV is $ENV"
 echo "INPUT_CSV is $INPUT_CSV"
-echo "OUTPUT_CSV is $OUTPUT_CSV"
 
 echo "Running script..."
 python data_tools/src/load_cans/main.py \
 --env "${ENV}" \
---input-csv "${INPUT_CSV}" \
---output-csv "${OUTPUT_CSV}"
+--input-csv "${INPUT_CSV}"

backend/data_tools/src/common/db.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from data_tools.environment.types import DataToolsConfig
+from sqlalchemy import Engine, create_engine
+from sqlalchemy.orm import Session, scoped_session, sessionmaker
+
+from models import *  # noqa: F403, F401
+from models import BaseModel
+from models.utils import track_db_history_after, track_db_history_before, track_db_history_catch_errors
+
+
+def init_db(
+    conn_string: str,
+) -> tuple[scoped_session[Session | Any], Engine]:  # noqa: F405
+    engine = create_engine(conn_string)
+
+    db_session = scoped_session(sessionmaker(autocommit=False, autoflush=False, bind=engine))
+
+    setup_schema(BaseModel)()
+
+    return db_session, engine
+
+def setup_triggers(session: scoped_session[Session | Any], sys_user: User) -> None:
+
+    @event.listens_for(session, "before_commit")
+    def receive_before_commit(session: Session):
+        track_db_history_before(session, sys_user)
+
+    @event.listens_for(session, "after_flush")
+    def receive_after_flush(session: Session, flush_context):
+        track_db_history_after(session, sys_user)
+
+    @event.listens_for(session.get_bind(), "handle_error")
+    def receive_error(exception_context):
+        track_db_history_catch_errors(exception_context)
+
+    return None
+
+
+def init_db_from_config(
+    config: DataToolsConfig, db: Optional[Engine] = None
+) -> tuple[sqlalchemy.engine.Engine, sqlalchemy.MetaData]:
+    if not db:
+        _, engine = init_db(config.db_connection_string)
+    else:
+        engine = db
+    return engine, BaseModel.metadata

backend/data_tools/src/common/utils.py

@@ -1,26 +1,19 @@
 from typing import Optional
+from uuid import UUID
 
-import sqlalchemy
 from data_tools.environment.azure import AzureConfig
 from data_tools.environment.dev import DevConfig
 from data_tools.environment.local import LocalConfig
 from data_tools.environment.pytest import PytestConfig
+from data_tools.environment.pytest_data_tools import PytestDataToolsConfig
 from data_tools.environment.types import DataToolsConfig
-from sqlalchemy import Engine, create_engine
+from sqlalchemy import select
+from sqlalchemy.orm import Session
 
-from models import BaseModel
+from models import User
 
-
-def init_db(
-    config: DataToolsConfig, db: Optional[Engine] = None
-) -> tuple[sqlalchemy.engine.Engine, sqlalchemy.MetaData]:
-    if not db:
-        engine = create_engine(
-            config.db_connection_string, echo=config.verbosity, future=True
-        )
-    else:
-        engine = db
-    return engine, BaseModel.metadata
+SYSTEM_ADMIN_OIDC_ID = "00000000-0000-1111-a111-000000000026"
+SYSTEM_ADMIN_EMAIL = "[email protected]"
 
 
 def get_config(environment_name: Optional[str] = None) -> DataToolsConfig:
@@ -33,6 +26,24 @@ def get_config(environment_name: Optional[str] = None) -> DataToolsConfig:
             config = LocalConfig()
         case "pytest":
             config = PytestConfig()
+        case "pytest_data_tools":
+            config = PytestDataToolsConfig()
         case _:
             config = DevConfig()
     return config
+
+def get_or_create_sys_user(session: Session) -> User:
+    """
+    Get or create the system user.
+
+    Args:
+        session: SQLAlchemy session object
+    Returns:
+        None
+    """
+    user = session.execute(select(User).where(User.oidc_id == SYSTEM_ADMIN_OIDC_ID)).scalar_one_or_none()
+
+    if not user:
+        user = User(email=SYSTEM_ADMIN_EMAIL, oidc_id=UUID(SYSTEM_ADMIN_OIDC_ID))
+
+    return user

backend/data_tools/src/disable_users/disable_users.py

@@ -1,23 +1,35 @@
-import logging
 import os
+import sys
+import time
+from datetime import timedelta
 
+from data_tools.src.common.db import init_db_from_config, setup_triggers
+from data_tools.src.common.utils import get_or_create_sys_user
 from data_tools.src.disable_users.queries import (
     ALL_ACTIVE_USER_SESSIONS_QUERY,
     EXCLUDED_USER_OIDC_IDS,
     GET_USER_ID_BY_OIDC_QUERY,
-    INACTIVE_USER_QUERY,
-    SYSTEM_ADMIN_EMAIL,
-    SYSTEM_ADMIN_OIDC_ID,
+    get_latest_user_session,
 )
 from data_tools.src.import_static_data.import_data import get_config, init_db
 from sqlalchemy import text
 from sqlalchemy.orm import Mapper, Session
 
 from models import *  # noqa: F403, F401
 
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+# Set the timezone to UTC
+os.environ["TZ"] = "UTC"
+time.tzset()
 
+# logger configuration
+format = (
+    "<green>{time:YYYY-MM-DD HH:mm:ss}</green> | "
+    "<level>{level: <8}</level> | "
+    "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> | "
+    "<level>{message}</level>"
+)
+logger.add(sys.stdout, format=format, level="INFO")
+logger.add(sys.stderr, format=format, level="INFO")
 
 def get_ids_from_oidc_ids(se, oidc_ids: list):
     """Retrieve user IDs corresponding to a list of OIDC IDs."""
@@ -33,46 +45,16 @@ def get_ids_from_oidc_ids(se, oidc_ids: list):
 
     return ids
 
-
-def create_system_admin(se):
-    """Create system user if it doesn't exist."""
-    system_admin = se.execute(
-        text(GET_USER_ID_BY_OIDC_QUERY),
-        {"oidc_id": SYSTEM_ADMIN_OIDC_ID}
-    ).fetchone()
-
-    if system_admin is None:
-        sys_user = User(
-            email=SYSTEM_ADMIN_EMAIL,
-            oidc_id=SYSTEM_ADMIN_OIDC_ID,
-            status=UserStatus.LOCKED
-        )
-        se.add(sys_user)
-        se.commit()
-        return sys_user.id
-
-    return system_admin[0]
-
-
 def disable_user(se, user_id, system_admin_id):
     """Deactivate a single user and log the change."""
     updated_user = User(id=user_id, status=UserStatus.INACTIVE, updated_by=system_admin_id)
     se.merge(updated_user)
 
-    db_audit = build_audit(updated_user, OpsDBHistoryType.UPDATED)
-    ops_db_history = OpsDBHistory(
-        event_type=OpsDBHistoryType.UPDATED,
-        created_by=system_admin_id,
-        class_name=updated_user.__class__.__name__,
-        row_key=db_audit.row_key,
-        changes=db_audit.changes,
-    )
-    se.add(ops_db_history)
-
     ops_event = OpsEvent(
         event_type=OpsEventType.UPDATE_USER,
         event_status=OpsEventStatus.SUCCESS,
         created_by=system_admin_id,
+        event_details={"user_id": user_id, "message": "User deactivated via automated process."},
     )
     se.add(ops_event)
 
@@ -88,12 +70,22 @@ def disable_user(se, user_id, system_admin_id):
 
 def update_disabled_users_status(conn: sqlalchemy.engine.Engine):
     """Update the status of disabled users in the database."""
+
     with Session(conn) as se:
         logger.info("Checking for System User.")
-        system_admin_id = create_system_admin(se)
+        system_admin = get_or_create_sys_user(se)
+        system_admin_id = system_admin.id
+
+        setup_triggers(se, system_admin)
 
         logger.info("Fetching inactive users.")
-        results = se.execute(text(INACTIVE_USER_QUERY)).scalars().all()
+        results = []
+        all_users = se.execute(select(User)).scalars().all()
+        for user in all_users:
+            latest_session = get_latest_user_session(user_id=user.id, session=se)
+            if latest_session and latest_session.last_active_at < datetime.now() - timedelta(days=60):
+                results.append(user.id)
+
         excluded_ids = get_ids_from_oidc_ids(se, EXCLUDED_USER_OIDC_IDS)
         user_ids = [uid for uid in results if uid not in excluded_ids]
 
@@ -115,9 +107,7 @@ def update_disabled_users_status(conn: sqlalchemy.engine.Engine):
 
     script_env = os.getenv("ENV")
     script_config = get_config(script_env)
-    db_engine, db_metadata_obj = init_db(script_config)
-
-    event.listen(Mapper, "after_configured", setup_schema(BaseModel))
+    db_engine, db_metadata_obj = init_db_from_config(script_config)
 
     update_disabled_users_status(db_engine)
 

backend/data_tools/src/disable_users/queries.py

@@ -1,5 +1,8 @@
-SYSTEM_ADMIN_OIDC_ID = "00000000-0000-1111-a111-000000000026"
-SYSTEM_ADMIN_EMAIL = "[email protected]"
+from data_tools.src.common.utils import SYSTEM_ADMIN_OIDC_ID
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+
+from models import UserSession
 
 EXCLUDED_USER_OIDC_IDS = [
     "00000000-0000-1111-a111-000000000018",     # Admin Demo
@@ -10,17 +13,6 @@
     SYSTEM_ADMIN_OIDC_ID                         # System Admin
 ]
 
-INACTIVE_USER_QUERY = (
-    "SELECT id "
-    "FROM ops_user "
-    "WHERE id IN ( "
-    "    SELECT ou.id "
-    "    FROM user_session JOIN ops_user ou ON user_session.user_id = ou.id "
-    "    WHERE ou.status = 'ACTIVE' "
-    "    AND user_session.last_active_at < CURRENT_TIMESTAMP - INTERVAL '60 days'"
-    ");"
-)
-
 ALL_ACTIVE_USER_SESSIONS_QUERY = (
     "SELECT * "
     "FROM user_session "
@@ -29,3 +21,14 @@
 )
 
 GET_USER_ID_BY_OIDC_QUERY = "SELECT id FROM ops_user WHERE oidc_id = :oidc_id"
+
+def get_latest_user_session(user_id: int, session: Session) -> UserSession | None:
+    return (
+        session.execute(
+            select(UserSession)
+            .where(UserSession.user_id == user_id)  # type: ignore
+            .order_by(UserSession.created_on.desc())
+        )
+        .scalars()
+        .first()
+    )

backend/data_tools/src/import_static_data/import_data.py

@@ -2,7 +2,8 @@
 import os
 
 import json5
-from data_tools.src.common.utils import get_config, init_db
+from data_tools.src.common.db import init_db, init_db_from_config
+from data_tools.src.common.utils import get_config
 from sqlalchemy import text
 from sqlalchemy.engine import Connection, Engine
 from sqlalchemy.orm import Session
@@ -157,7 +158,7 @@ def import_data(engine: Engine, data: dict[str, Any]) -> None:
     script_env = os.getenv("ENV")
     script_config = get_config(script_env)
 
-    db_engine, db_metadata_obj = init_db(script_config)
+    db_engine, db_metadata_obj = init_db_from_config(script_config)
 
     global_data = get_data_to_import()
 

backend/data_tools/src/import_static_data/load_db.py

@@ -1,7 +1,8 @@
 import os
 
 import sqlalchemy.engine
-from data_tools.src.common.utils import get_config, init_db
+from data_tools.src.common.db import init_db_from_config
+from data_tools.src.common.utils import get_config
 from sqlalchemy.orm import configure_mappers
 
 from models import BaseModel
@@ -18,7 +19,7 @@ def delete_and_create(engine: sqlalchemy.engine.Engine) -> None:
     script_env = os.getenv("ENV")
     script_config = get_config(script_env)
 
-    db_engine, db_metadata_obj = init_db(script_config)
+    db_engine, db_metadata_obj = init_db_from_config(script_config)
 
     delete_and_create(db_engine)