1st iter

backend/src/core/database.py

@@ -12,8 +12,9 @@
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker, Session
 from backend.src.models.mapping import Base
-# Import TaskRecord to ensure it's registered with Base
+# Import models to ensure they're registered with Base
 from backend.src.models.task import TaskRecord
+from backend.src.models.connection import ConnectionConfig
 import os
 # [/SECTION]
 

backend/src/core/superset_client.py

@@ -78,6 +78,28 @@ class SupersetClient(BaseSupersetClient):
         return result
     # [/DEF:SupersetClient.get_dashboards_summary:Function]
 
+    # [DEF:SupersetClient.get_dataset:Function]
+    # @PURPOSE: Fetch full dataset structure including columns and metrics.
+    # @PARAM:   dataset_id (int) - The ID of the dataset.
+    # @RETURN:  Dict - The dataset metadata.
+    def get_dataset(self, dataset_id: int) -> Dict:
+        """
+        Fetch full dataset structure.
+        """
+        return self.network.get(f"/api/v1/dataset/{dataset_id}").json()
+    # [/DEF:SupersetClient.get_dataset:Function]
+
+    # [DEF:SupersetClient.update_dataset:Function]
+    # @PURPOSE: Update dataset metadata.
+    # @PARAM:   dataset_id (int) - The ID of the dataset.
+    # @PARAM:   data (Dict) - The payload for update.
+    def update_dataset(self, dataset_id: int, data: Dict):
+        """
+        Update dataset metadata.
+        """
+        self.network.put(f"/api/v1/dataset/{dataset_id}", json=data)
+    # [/DEF:SupersetClient.update_dataset:Function]
+
 # [/DEF:SupersetClient:Class]
 
 # [/DEF:backend.src.core.superset_client:Module]

backend/src/core/task_manager/manager.py

@@ -98,9 +98,9 @@ class TaskManager:
                 params = {**task.params, "_task_id": task_id}
                 
                 if asyncio.iscoroutinefunction(plugin.execute):
-                    await plugin.execute(params)
+                    task.result = await plugin.execute(params)
                 else:
-                    await self.loop.run_in_executor(
+                    task.result = await self.loop.run_in_executor(
                         self.executor,
                         plugin.execute,
                         params

backend/src/core/task_manager/models.py

@@ -51,6 +51,7 @@ class Task(BaseModel):
     params: Dict[str, Any] = Field(default_factory=dict)
     input_required: bool = False
     input_request: Optional[Dict[str, Any]] = None
+    result: Optional[Dict[str, Any]] = None
 
     # [DEF:Task.__init__:Function]
     # @PURPOSE: Initializes the Task model and validates input_request for AWAITING_INPUT status.

backend/src/core/task_manager/persistence.py

@@ -43,6 +43,7 @@ class TaskPersistenceService:
                 record.started_at = task.started_at
                 record.finished_at = task.finished_at
                 record.params = task.params
+                record.result = task.result
                 
                 # Store logs as JSON, converting datetime to string
                 record.logs = []
@@ -108,6 +109,7 @@ class TaskPersistenceService:
                             started_at=record.started_at,
                             finished_at=record.finished_at,
                             params=record.params or {},
+                            result=record.result,
                             logs=logs
                         )
                         loaded_tasks.append(task)