# #region DatasetsApi [C:3] [TYPE Module] [SEMANTICS fastapi, dataset, api, search, mapping, mapped-fields] # # @BRIEF API endpoints for the Dataset Hub - listing datasets with mapping progress # @LAYER: API # @RELATION DEPENDS_ON -> [AppDependencies] # @RELATION DEPENDS_ON -> [ResourceService] # @RELATION DEPENDS_ON -> [SupersetClient] # # @INVARIANT: All dataset responses include last_task metadata from fastapi import APIRouter, Depends, HTTPException from pydantic import BaseModel, Field from ...core.logger import belief_scope, logger from ...core.superset_client import SupersetClient from ...dependencies import get_config_manager, get_resource_service, get_task_manager, has_permission router = APIRouter(prefix="/api/datasets", tags=["Datasets"]) # #region MappedFields [C:1] [TYPE DataClass] # @BRIEF DTO for dataset mapping progress statistics class MappedFields(BaseModel): total: int mapped: int # #endregion MappedFields # #region LastTask [C:1] [TYPE DataClass] # @BRIEF DTO for the most recent task associated with a dataset class LastTask(BaseModel): task_id: str | None = None status: str | None = Field(None, pattern="^RUNNING|SUCCESS|ERROR|WAITING_INPUT$") # #endregion LastTask # #region DatasetItem [C:1] [TYPE DataClass] # @BRIEF Summary DTO for a dataset in the hub listing class DatasetItem(BaseModel): id: int table_name: str schema_name: str = Field(..., alias="schema") database: str mapped_fields: MappedFields | None = None last_task: LastTask | None = None class Config: allow_population_by_field_name = True # #endregion DatasetItem # #region LinkedDashboard [C:1] [TYPE DataClass] # @BRIEF DTO for a dashboard linked to a dataset class LinkedDashboard(BaseModel): id: int title: str slug: str | None = None # #endregion LinkedDashboard # #region DatasetColumn [C:1] [TYPE DataClass] # @BRIEF DTO for a single dataset column's metadata class DatasetColumn(BaseModel): id: int name: str type: str | None = None is_dttm: bool = False is_active: bool = True description: str | None = None # #endregion DatasetColumn # #region DatasetDetailResponse [C:1] [TYPE DataClass] # @BRIEF Detailed DTO for a dataset including columns and links class DatasetDetailResponse(BaseModel): id: int table_name: str | None = None schema_name: str | None = Field(None, alias="schema") database: str description: str | None = None columns: list[DatasetColumn] column_count: int sql: str | None = None linked_dashboards: list[LinkedDashboard] linked_dashboard_count: int is_sqllab_view: bool = False created_on: str | None = None changed_on: str | None = None class Config: allow_population_by_field_name = True # #endregion DatasetDetailResponse # #region DatasetsResponse [C:1] [TYPE DataClass] # @BRIEF Paginated response DTO for dataset listings class DatasetsResponse(BaseModel): datasets: list[DatasetItem] total: int page: int page_size: int total_pages: int # #endregion DatasetsResponse # #region TaskResponse [C:1] [TYPE DataClass] # @BRIEF Response DTO containing a task ID for tracking class TaskResponse(BaseModel): task_id: str # #endregion TaskResponse # #region get_dataset_ids [C:3] [TYPE Function] # @BRIEF Fetch list of all dataset IDs from a specific environment (without pagination) # @PRE: env_id must be a valid environment ID # @POST: Returns a list of all dataset IDs # @RELATION CALLS -> [get_datasets_with_status] @router.get("/ids") async def get_dataset_ids( env_id: str, search: str | None = None, config_manager=Depends(get_config_manager), task_manager=Depends(get_task_manager), resource_service=Depends(get_resource_service), _ = Depends(has_permission("plugin:migration", "READ")) ): with belief_scope("get_dataset_ids", f"env_id={env_id}, search={search}"): # Validate environment exists environments = config_manager.get_environments() env = next((e for e in environments if e.id == env_id), None) if not env: logger.error(f"[get_dataset_ids][Coherence:Failed] Environment not found: {env_id}") raise HTTPException(status_code=404, detail="Environment not found") try: # Get all tasks for status lookup all_tasks = task_manager.get_all_tasks() # Fetch datasets with status using ResourceService datasets = await resource_service.get_datasets_with_status(env, all_tasks) # Apply search filter if provided if search: search_lower = search.lower() datasets = [ d for d in datasets if search_lower in d.get('table_name', '').lower() ] # Extract and return just the IDs dataset_ids = [d['id'] for d in datasets] logger.info(f"[get_dataset_ids][Coherence:OK] Returning {len(dataset_ids)} dataset IDs") return {"dataset_ids": dataset_ids} except Exception as e: logger.error(f"[get_dataset_ids][Coherence:Failed] Failed to fetch dataset IDs: {e}") raise HTTPException(status_code=503, detail=f"Failed to fetch dataset IDs: {e!s}") # #endregion get_dataset_ids # #region get_datasets [C:3] [TYPE Function] # @BRIEF Fetch list of datasets from a specific environment with mapping progress # @PRE: env_id must be a valid environment ID # @PRE: page must be >= 1 if provided # @PRE: page_size must be between 1 and 100 if provided # @POST: Returns a list of datasets with enhanced metadata and pagination info # @POST: Response includes pagination metadata (page, page_size, total, total_pages) # @RELATION CALLS -> [get_datasets_with_status] @router.get("", response_model=DatasetsResponse) async def get_datasets( env_id: str, search: str | None = None, page: int = 1, page_size: int = 10, config_manager=Depends(get_config_manager), task_manager=Depends(get_task_manager), resource_service=Depends(get_resource_service), _ = Depends(has_permission("plugin:migration", "READ")) ): with belief_scope("get_datasets", f"env_id={env_id}, search={search}, page={page}, page_size={page_size}"): # Validate pagination parameters if page < 1: logger.error(f"[get_datasets][Coherence:Failed] Invalid page: {page}") raise HTTPException(status_code=400, detail="Page must be >= 1") if page_size < 1 or page_size > 100: logger.error(f"[get_datasets][Coherence:Failed] Invalid page_size: {page_size}") raise HTTPException(status_code=400, detail="Page size must be between 1 and 100") # Validate environment exists environments = config_manager.get_environments() env = next((e for e in environments if e.id == env_id), None) if not env: logger.error(f"[get_datasets][Coherence:Failed] Environment not found: {env_id}") raise HTTPException(status_code=404, detail="Environment not found") try: # Get all tasks for status lookup all_tasks = task_manager.get_all_tasks() # Fetch datasets with status using ResourceService datasets = await resource_service.get_datasets_with_status(env, all_tasks) # Apply search filter if provided if search: search_lower = search.lower() datasets = [ d for d in datasets if search_lower in d.get('table_name', '').lower() ] # Calculate pagination total = len(datasets) total_pages = (total + page_size - 1) // page_size if total > 0 else 1 start_idx = (page - 1) * page_size end_idx = start_idx + page_size # Slice datasets for current page paginated_datasets = datasets[start_idx:end_idx] logger.info(f"[get_datasets][Coherence:OK] Returning {len(paginated_datasets)} datasets (page {page}/{total_pages}, total: {total})") return DatasetsResponse( datasets=paginated_datasets, total=total, page=page, page_size=page_size, total_pages=total_pages ) except Exception as e: logger.error(f"[get_datasets][Coherence:Failed] Failed to fetch datasets: {e}") raise HTTPException(status_code=503, detail=f"Failed to fetch datasets: {e!s}") # #endregion get_datasets # #region MapColumnsRequest [C:1] [TYPE DataClass] # @BRIEF Request DTO for initiating column mapping class MapColumnsRequest(BaseModel): env_id: str = Field(..., description="Environment ID") dataset_ids: list[int] = Field(..., description="List of dataset IDs to map") source_type: str = Field(..., description="Source type: 'postgresql' or 'xlsx'") connection_id: str | None = Field(None, description="Connection ID for PostgreSQL source") file_data: str | None = Field(None, description="File path or data for XLSX source") # #endregion MapColumnsRequest # #region map_columns [C:3] [TYPE Function] # @BRIEF Trigger bulk column mapping for datasets # @PRE: User has permission plugin:mapper:execute # @PRE: env_id is a valid environment ID # @PRE: dataset_ids is a non-empty list # @POST: Returns task_id for tracking mapping progress # @POST: Task is created and queued for execution # @RELATION DISPATCHES -> [MapperPlugin] # @RELATION CALLS -> [create_task] @router.post("/map-columns", response_model=TaskResponse) async def map_columns( request: MapColumnsRequest, config_manager=Depends(get_config_manager), task_manager=Depends(get_task_manager), _ = Depends(has_permission("plugin:mapper", "EXECUTE")) ): with belief_scope("map_columns", f"env={request.env_id}, count={len(request.dataset_ids)}, source={request.source_type}"): # Validate request if not request.dataset_ids: logger.error("[map_columns][Coherence:Failed] No dataset IDs provided") raise HTTPException(status_code=400, detail="At least one dataset ID must be provided") # Validate source type if request.source_type not in ['postgresql', 'xlsx']: logger.error(f"[map_columns][Coherence:Failed] Invalid source type: {request.source_type}") raise HTTPException(status_code=400, detail="Source type must be 'postgresql' or 'xlsx'") # Validate environment exists environments = config_manager.get_environments() env = next((e for e in environments if e.id == request.env_id), None) if not env: logger.error(f"[map_columns][Coherence:Failed] Environment not found: {request.env_id}") raise HTTPException(status_code=404, detail="Environment not found") try: # Create mapping task task_params = { 'env': request.env_id, 'dataset_id': request.dataset_ids[0] if request.dataset_ids else None, 'source': request.source_type, 'connection_id': request.connection_id, 'file_data': request.file_data } task_obj = await task_manager.create_task( plugin_id='dataset-mapper', params=task_params ) logger.info(f"[map_columns][Coherence:OK] Mapping task created: {task_obj.id} for {len(request.dataset_ids)} datasets") return TaskResponse(task_id=str(task_obj.id)) except Exception as e: logger.error(f"[map_columns][Coherence:Failed] Failed to create mapping task: {e}") raise HTTPException(status_code=503, detail=f"Failed to create mapping task: {e!s}") # #endregion map_columns # #region GenerateDocsRequest [C:1] [TYPE DataClass] # @BRIEF Request DTO for initiating documentation generation class GenerateDocsRequest(BaseModel): env_id: str = Field(..., description="Environment ID") dataset_ids: list[int] = Field(..., description="List of dataset IDs to generate docs for") llm_provider: str = Field(..., description="LLM provider to use") options: dict | None = Field(None, description="Additional options for documentation generation") # #endregion GenerateDocsRequest # #region generate_docs [C:3] [TYPE Function] # @BRIEF Trigger bulk documentation generation for datasets # @PRE: User has permission plugin:llm_analysis:execute # @PRE: env_id is a valid environment ID # @PRE: dataset_ids is a non-empty list # @POST: Returns task_id for tracking documentation generation progress # @POST: Task is created and queued for execution # @RELATION DISPATCHES -> [DocumentationPlugin] # @RELATION CALLS -> [create_task] @router.post("/generate-docs", response_model=TaskResponse) async def generate_docs( request: GenerateDocsRequest, config_manager=Depends(get_config_manager), task_manager=Depends(get_task_manager), _ = Depends(has_permission("plugin:llm_analysis", "EXECUTE")) ): with belief_scope("generate_docs", f"env={request.env_id}, count={len(request.dataset_ids)}, provider={request.llm_provider}"): # Validate request if not request.dataset_ids: logger.error("[generate_docs][Coherence:Failed] No dataset IDs provided") raise HTTPException(status_code=400, detail="At least one dataset ID must be provided") # Validate environment exists environments = config_manager.get_environments() env = next((e for e in environments if e.id == request.env_id), None) if not env: logger.error(f"[generate_docs][Coherence:Failed] Environment not found: {request.env_id}") raise HTTPException(status_code=404, detail="Environment not found") try: # Create documentation generation task task_params = { 'environment_id': request.env_id, 'dataset_id': str(request.dataset_ids[0]) if request.dataset_ids else None, 'provider_id': request.llm_provider, 'options': request.options or {} } task_obj = await task_manager.create_task( plugin_id='llm_documentation', params=task_params ) logger.info(f"[generate_docs][Coherence:OK] Documentation generation task created: {task_obj.id} for {len(request.dataset_ids)} datasets") return TaskResponse(task_id=str(task_obj.id)) except Exception as e: logger.error(f"[generate_docs][Coherence:Failed] Failed to create documentation generation task: {e}") raise HTTPException(status_code=503, detail=f"Failed to create documentation generation task: {e!s}") # #endregion generate_docs # #region get_dataset_detail [C:3] [TYPE Function] # @BRIEF Get detailed dataset information including columns and linked dashboards # @PRE: env_id is a valid environment ID # @PRE: dataset_id is a valid dataset ID # @POST: Returns detailed dataset info with columns and linked dashboards # @RELATION CALLS -> [SupersetClientGetDatasetDetail] @router.get("/{dataset_id}", response_model=DatasetDetailResponse) async def get_dataset_detail( env_id: str, dataset_id: int, config_manager=Depends(get_config_manager), _ = Depends(has_permission("plugin:migration", "READ")) ): with belief_scope("get_dataset_detail", f"env_id={env_id}, dataset_id={dataset_id}"): # Validate environment exists environments = config_manager.get_environments() env = next((e for e in environments if e.id == env_id), None) if not env: logger.error(f"[get_dataset_detail][Coherence:Failed] Environment not found: {env_id}") raise HTTPException(status_code=404, detail="Environment not found") try: # Fetch detailed dataset info using SupersetClient client = SupersetClient(env) dataset_detail = client.get_dataset_detail(dataset_id) logger.info(f"[get_dataset_detail][Coherence:OK] Retrieved dataset {dataset_id} with {dataset_detail['column_count']} columns and {dataset_detail['linked_dashboard_count']} linked dashboards") return DatasetDetailResponse(**dataset_detail) except Exception as e: logger.error(f"[get_dataset_detail][Coherence:Failed] Failed to fetch dataset detail: {e}") raise HTTPException(status_code=503, detail=f"Failed to fetch dataset detail: {e!s}") # #endregion get_dataset_detail # #endregion DatasetsApi