# #region IdMappingServiceModule [C:5] [TYPE Module] [SEMANTICS sqlalchemy, mapping, sync, schedule, superset, resource] # @defgroup Core Module group. # # @BRIEF Service for tracking and synchronizing Superset Resource IDs (UUID <-> Integer ID) # @LAYER Core # @RELATION DEPENDS_ON -> [MappingModels] # @RELATION DEPENDS_ON -> [LoggerModule] # @PRE Database session is valid and Superset client factory returns authenticated clients for requested environments. # @POST Mapping synchronization and lookup APIs are available for environment-scoped UUID-to-integer resolution. # @SIDE_EFFECT Reads/writes ResourceMapping rows, emits logs, and schedules periodic sync jobs. # @DATA_CONTRACT Input[environment_id, resource_type, uuid] -> Output[remote_integer_id|None] # @TEST_DATA: mock_superset_resources -> {'chart': [{'id': 42, 'uuid': '1234', 'slice_name': 'test'}], 'dataset': [{'id': 99, 'uuid': '5678', 'table_name': 'data'}]} # # @INVARIANT sync_environment must handle remote API failures gracefully. from datetime import UTC, datetime from apscheduler.schedulers.background import BackgroundScheduler from apscheduler.triggers.cron import CronTrigger from sqlalchemy.orm import Session from src.core.cot_logger import seed_trace_id from src.core.logger import belief_scope, logger from src.models.mapping import ResourceMapping, ResourceType # #region IdMappingService [C:5] [TYPE Class] # @defgroup Core Module group. # @BRIEF Service handling the cataloging and retrieval of remote Superset Integer IDs. # @PRE db_session is an active SQLAlchemy Session bound to mapping tables. # @POST Service instance provides scheduler control and environment-scoped mapping synchronization APIs. # @RELATION DEPENDS_ON -> [MappingModels] # @RELATION DEPENDS_ON -> [LoggerModule] # @INVARIANT self.db remains the authoritative session for all mapping operations. # @SIDE_EFFECT Instantiates an in-process scheduler and performs database writes during sync cycles. # @DATA_CONTRACT Input[db_session] -> Output[IdMappingService] # # @TEST_CONTRACT IdMappingServiceModel -> # { # required_fields: {db_session: Session}, # invariants: [ # "sync_environment correctly creates or updates ResourceMapping records", # "get_remote_id returns an integer or None", # "get_remote_ids_batch returns a dictionary of valid UUIDs to integers" # ] # } # @TEST_FIXTURE valid_mapping_service -> {"db_session": "MockSession()"} # @TEST_EDGE sync_api_failure -> handles exception gracefully # @TEST_EDGE get_remote_id_not_found -> returns None # @TEST_EDGE get_batch_empty_list -> returns empty dict # @TEST_INVARIANT resilient_fetching -> verifies: [sync_api_failure] class IdMappingService: # #region __init__ [TYPE Function] # @PURPOSE: Initializes the mapping service. def __init__(self, db_session: Session): self.db = db_session self.scheduler = BackgroundScheduler() self._sync_job = None # #endregion __init__ # #region start_scheduler [TYPE Function] # @ingroup Core # @PURPOSE: Starts the background scheduler with a given cron string. # @PARAM cron_string (str) - Cron expression for the sync interval. # @PARAM environments (List[str]) - List of environment IDs to sync. # @PARAM superset_client_factory - Function to get a client for an environment. def start_scheduler( self, cron_string: str, environments: list[str], superset_client_factory ): with belief_scope("IdMappingService.start_scheduler"): if self._sync_job: self.scheduler.remove_job(self._sync_job.id) logger.info( "[IdMappingService.start_scheduler][Reflect] Removed existing sync job." ) def sync_all(): seed_trace_id() for env_id in environments: client = superset_client_factory(env_id) if client: import asyncio asyncio.run(self.sync_environment(env_id, client)) self._sync_job = self.scheduler.add_job( sync_all, CronTrigger.from_crontab(cron_string), id="id_mapping_sync_job", replace_existing=True, ) if not self.scheduler.running: self.scheduler.start() logger.info( f"[IdMappingService.start_scheduler][Coherence:OK] Started background scheduler with cron: {cron_string}" ) else: logger.info( f"[IdMappingService.start_scheduler][Coherence:OK] Updated background scheduler with cron: {cron_string}" ) # #endregion start_scheduler # #region sync_environment [TYPE Function] # @ingroup Core # @PURPOSE: Fully synchronizes mapping for a specific environment. # @PARAM environment_id (str) - Target environment ID. # @PARAM superset_client - Instance capable of hitting the Superset API. # @PRE environment_id exists in the database. # @POST ResourceMapping records for the environment are created or updated. async def sync_environment( self, environment_id: str, superset_client, incremental: bool = False ) -> None: """ Polls the Superset APIs for the target environment and updates the local mapping table. If incremental=True, only fetches items changed since the max last_synced_at date. """ seed_trace_id() with belief_scope("IdMappingService.sync_environment"): logger.reason( f"Starting sync for environment {environment_id}", extra={"src": "IdMappingService.sync_environment", "payload": {"environment_id": environment_id, "incremental": incremental}}, ) # Implementation Note: In a real scenario, superset_client needs to be an instance # capable of auth & iteration over /api/v1/chart/, /api/v1/dataset/, /api/v1/dashboard/ # Here we structure the logic according to the spec. types_to_poll = [ (ResourceType.CHART, "chart", "slice_name"), (ResourceType.DATASET, "dataset", "table_name"), ( ResourceType.DASHBOARD, "dashboard", "slug", ), # Note: dashboard slug or dashboard_title ] total_synced = 0 total_deleted = 0 try: for res_enum, endpoint, name_field in types_to_poll: logger.debug( f"[IdMappingService.sync_environment][Explore] Polling {endpoint} endpoint" ) # Simulated API Fetch (Would be: superset_client.get(f"/api/v1/{endpoint}/")... ) # This relies on the superset API structure, e.g. { "result": [{"id": 1, "uuid": "...", name_field: "..."}] } # We assume superset_client provides a generic method to fetch all pages. try: since_dttm = None if incremental: from sqlalchemy.sql import func max_date = ( self.db.query(func.max(ResourceMapping.last_synced_at)) .filter( ResourceMapping.environment_id == environment_id, ResourceMapping.resource_type == res_enum, ) .scalar() ) if max_date: # We subtract a bit for safety overlap from datetime import timedelta since_dttm = max_date - timedelta(minutes=5) logger.debug( f"[IdMappingService.sync_environment] Incremental sync since {since_dttm}" ) resources = await superset_client.get_all_resources( endpoint, since_dttm=since_dttm ) # Track which UUIDs we see in this sync cycle synced_uuids = set() for res in resources: res_uuid = res.get("uuid") raw_id = res.get("id") res_name = res.get(name_field) if not res_uuid or raw_id is None: continue synced_uuids.add(res_uuid) res_id = str(raw_id) # Store as string # Upsert Logic mapping = ( self.db.query(ResourceMapping) .filter_by( environment_id=environment_id, resource_type=res_enum, uuid=res_uuid, ) .first() ) if mapping: mapping.remote_integer_id = res_id mapping.resource_name = res_name mapping.last_synced_at = datetime.now(UTC) else: new_mapping = ResourceMapping( environment_id=environment_id, resource_type=res_enum, uuid=res_uuid, remote_integer_id=res_id, resource_name=res_name, last_synced_at=datetime.now(UTC), ) self.db.add(new_mapping) total_synced += 1 # Delete stale mappings: rows for this env+type whose UUID # was NOT returned by the API (resource was deleted remotely) # We only do this on full syncs, because incremental syncs don't return all UUIDs if not incremental: stale_query = self.db.query(ResourceMapping).filter( ResourceMapping.environment_id == environment_id, ResourceMapping.resource_type == res_enum, ) if synced_uuids: stale_query = stale_query.filter( ResourceMapping.uuid.notin_(synced_uuids) ) deleted = stale_query.delete(synchronize_session="fetch") if deleted: total_deleted += deleted logger.reason( f"Removed {deleted} stale {endpoint} mapping(s) for {environment_id}", extra={"src": "IdMappingService.sync_environment", "payload": {"deleted": deleted, "endpoint": endpoint, "environment_id": environment_id}}, ) except Exception as loop_e: logger.error( f"[IdMappingService.sync_environment][Reason] Error polling {endpoint}: {loop_e}" ) # Continue to next resource type instead of blowing up the whole sync self.db.commit() logger.info( f"[IdMappingService.sync_environment][Coherence:OK] Successfully synced {total_synced} items and deleted {total_deleted} stale items." ) except Exception as e: self.db.rollback() logger.error( f"[IdMappingService.sync_environment][Coherence:Failed] Critical sync failure: {e}" ) raise # #endregion sync_environment # #region get_remote_id [TYPE Function] # @ingroup Core # @PURPOSE: Retrieves the remote integer ID for a given universal UUID. # @PARAM environment_id (str) # @PARAM resource_type (ResourceType) # @PARAM uuid (str) # @RETURN Optional[int] def get_remote_id( self, environment_id: str, resource_type: ResourceType, uuid: str ) -> int | None: mapping = ( self.db.query(ResourceMapping) .filter_by( environment_id=environment_id, resource_type=resource_type, uuid=uuid ) .first() ) if mapping: try: return int(mapping.remote_integer_id) except ValueError: return None return None # #endregion get_remote_id # #region get_remote_ids_batch [TYPE Function] # @ingroup Core # @PURPOSE: Retrieves remote integer IDs for a list of universal UUIDs efficiently. # @PARAM environment_id (str) # @PARAM resource_type (ResourceType) # @PARAM uuids (List[str]) # @RETURN Dict[str, int] - Mapping of UUID -> Integer ID def get_remote_ids_batch( self, environment_id: str, resource_type: ResourceType, uuids: list[str] ) -> dict[str, int]: if not uuids: return {} mappings = ( self.db.query(ResourceMapping) .filter( ResourceMapping.environment_id == environment_id, ResourceMapping.resource_type == resource_type, ResourceMapping.uuid.in_(uuids), ) .all() ) result = {} for m in mappings: try: result[m.uuid] = int(m.remote_integer_id) except ValueError: logger.debug("Could not parse remote_integer_id for mapping %s (uuid=%s)", m.id, m.uuid) return result # #endregion get_remote_ids_batch # #endregion IdMappingService # #endregion IdMappingServiceModule