diff --git a/app/db/models.py b/app/db/models.py
index 9b4331a3cf27d4cba4936fc2342546b441c3047a..07183a3e0c5c3c5bc1596f9015df06c90ed506fc 100644
--- a/app/db/models.py
+++ b/app/db/models.py
@@ -4,6 +4,7 @@ from uuid import UUID
import pydantic
from mitm_tooling.representation import Header, SQLRepresentationSchema, mk_sql_rep_schema
from mitm_tooling.transformation.superset.common import SupersetDBConnectionInfo
+from mitm_tooling.transformation.superset.definitions.mitm_dataset import MitMDatasetIdentifier
from pydantic import AnyUrl, BaseModel
from sqlmodel import SQLModel, Field
@@ -33,6 +34,10 @@ class TrackedMitMDataset(SQLModel, table=True):
sql_alchemy_uri: AnyUrl = Field(sa_type=StrType.wrap(AnyUrl))
mitm_header: Header = Field(sa_type=PydanticType.wrap(Header), repr=False)
+ @property
+ def identifier(self) -> MitMDatasetIdentifier:
+ return MitMDatasetIdentifier(dataset_name=self.dataset_name, uuid=self.uuid)
+
@property
def superset_connection_info(self) -> SupersetDBConnectionInfo:
return SupersetDBConnectionInfo(sql_alchemy_uri=self.sql_alchemy_uri, schema_name=self.schema_name)
diff --git a/app/routes/definitions/requests.py b/app/routes/definitions/requests.py
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..08cf67ebb9afa5d1e5285d27e97cfa750d47d0e3 100644
--- a/app/routes/definitions/requests.py
+++ b/app/routes/definitions/requests.py
@@ -0,0 +1,10 @@
+import pydantic
+from mitm_tooling.representation import Header
+from mitm_tooling.transformation.superset.common import SupersetDBConnectionInfo
+from mitm_tooling.transformation.superset.definitions.mitm_dataset import MitMDatasetIdentifier
+
+
+class GenerateMitMDatasetDefinitionRequest(pydantic.BaseModel):
+ mitm_header: Header
+ dataset_identifier: MitMDatasetIdentifier
+ db_conn_info: SupersetDBConnectionInfo
diff --git a/app/routes/definitions/responses.py b/app/routes/definitions/responses.py
index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..1869bc23da2f9aa5918a2bd468c62c110e5309df 100644
--- a/app/routes/definitions/responses.py
+++ b/app/routes/definitions/responses.py
@@ -0,0 +1,8 @@
+from mitm_tooling.transformation.superset.definitions import SupersetMitMDatasetImport
+from mitm_tooling.transformation.superset.definition_bundles import SupersetMitMDatasetBundle
+
+class MitMDatasetBundleResponse(SupersetMitMDatasetBundle):
+ pass
+
+class MitMDatasetImportResponse(SupersetMitMDatasetImport):
+ pass
\ No newline at end of file
diff --git a/app/routes/definitions/router.py b/app/routes/definitions/router.py
index b169db7f0a6d8611cc0253e2a01d06cd056d95a2..ee649798be69d492f6579b7aa6c44a3347e450c6 100644
--- a/app/routes/definitions/router.py
+++ b/app/routes/definitions/router.py
@@ -2,77 +2,74 @@ import io
import logging
from fastapi import APIRouter
-from mitm_tooling.representation import Header
from mitm_tooling.transformation.superset import write_superset_import_as_zip
-from mitm_tooling.transformation.superset.common import SupersetDBConnectionInfo
-from mitm_tooling.transformation.superset.definition_bundles import SupersetMitMDatasetBundle
from starlette.responses import StreamingResponse
from app.dependencies.orm import TrackedMitMDatasetDependency
+from app.routes.definitions.requests import GenerateMitMDatasetDefinitionRequest
+from app.routes.definitions.responses import MitMDatasetBundleResponse, MitMDatasetImportResponse
router = APIRouter(prefix='/definitions', tags=['definitions'])
logger = logging.getLogger(__name__)
-from mitm_tooling.transformation.superset.definitions import SupersetMitMDatasetImport
from mitm_tooling.transformation.superset.interface import mk_superset_mitm_dataset_bundle, \
mk_superset_visualization_bundle
@router.post('/mitm_dataset')
-def generate_mitm_dataset_bundle(dataset_name: str,
- mitm_header: Header,
- db_conn_info: SupersetDBConnectionInfo,
- include_visualizations: bool = False) -> SupersetMitMDatasetBundle:
- mitm_dataset_bundle = mk_superset_mitm_dataset_bundle(dataset_name, mitm_header, db_conn_info)
+def generate_mitm_dataset_bundle(request: GenerateMitMDatasetDefinitionRequest,
+ include_visualizations: bool = False) -> MitMDatasetBundleResponse:
+ mitm_dataset_bundle = mk_superset_mitm_dataset_bundle(request.mitm_header, request.dataset_identifier,
+ request.db_conn_info)
if include_visualizations:
- mitm_dataset_bundle.visualization_bundle = mk_superset_visualization_bundle(mitm_header,
+ mitm_dataset_bundle.visualization_bundle = mk_superset_visualization_bundle(request.mitm_header,
mitm_dataset_bundle.datasource_bundle)
- return mitm_dataset_bundle
+ return MitMDatasetBundleResponse(**mitm_dataset_bundle.model_dump())
@router.post('/mitm_dataset/import')
-def generate_mitm_dataset_import(dataset_name: str,
- mitm_header: Header,
- db_conn_info: SupersetDBConnectionInfo,
- include_visualizations: bool = False) -> SupersetMitMDatasetImport:
- return generate_mitm_dataset_bundle(dataset_name, mitm_header, db_conn_info, include_visualizations).to_import()
+def generate_mitm_dataset_import(request: GenerateMitMDatasetDefinitionRequest,
+ include_visualizations: bool = False) -> MitMDatasetImportResponse:
+ return MitMDatasetImportResponse(
+ **generate_mitm_dataset_bundle(request, include_visualizations).to_import().model_dump())
@router.post('/mitm_dataset/import/zip', response_class=StreamingResponse,
responses={200: {'content': {'application/zip': {}}}})
-def generate_mitm_dataset_import_zip(dataset_name: str,
- mitm_header: Header,
- db_conn_info: SupersetDBConnectionInfo,
+def generate_mitm_dataset_import_zip(request: GenerateMitMDatasetDefinitionRequest,
include_visualizations: bool = False) -> StreamingResponse:
- assets = generate_mitm_dataset_bundle(dataset_name, mitm_header, db_conn_info, include_visualizations)
+ assets = generate_mitm_dataset_bundle(request, include_visualizations)
bio = io.BytesIO()
write_superset_import_as_zip(bio, assets)
bio.seek(0)
return StreamingResponse(bio, media_type='application/zip')
-@router.get('/mitm_dataset/{uuid}/')
-def generate_uploaded_mitm_dataset_bundle(uploaded_dataset: TrackedMitMDatasetDependency,
- include_visualizations: bool = False) -> SupersetMitMDatasetBundle:
- return generate_mitm_dataset_bundle(uploaded_dataset.dataset_name, uploaded_dataset.mitm_header,
- uploaded_dataset.superset_connection_info,
- include_visualizations=include_visualizations)
+@router.get('/mitm_dataset/{uuid}')
+def generate_tracked_mitm_dataset_bundle(tracked_dataset: TrackedMitMDatasetDependency,
+ include_visualizations: bool = False) -> MitMDatasetBundleResponse:
+ request = GenerateMitMDatasetDefinitionRequest(mitm_header=tracked_dataset.mitm_header,
+ dataset_identifier=tracked_dataset.dataset_identifier,
+ db_conn_info=tracked_dataset.db_conn_info)
+ return MitMDatasetBundleResponse(
+ **generate_mitm_dataset_bundle(request,
+ include_visualizations=include_visualizations).model_dump())
@router.get('/mitm_dataset/{uuid}/import')
-def generate_uploaded_mitm_dataset_import(uploaded_dataset: TrackedMitMDatasetDependency,
- include_visualizations: bool = False) -> SupersetMitMDatasetImport:
- mitm_dataset_bundle = generate_uploaded_mitm_dataset_bundle(uploaded_dataset,
- include_visualizations=include_visualizations)
- return mitm_dataset_bundle.to_import()
+def generate_tracked_mitm_dataset_import(tracked_dataset: TrackedMitMDatasetDependency,
+ include_visualizations: bool = False) -> MitMDatasetImportResponse:
+ mitm_dataset_bundle = generate_tracked_mitm_dataset_bundle(tracked_dataset,
+ include_visualizations=include_visualizations)
+ return MitMDatasetImportResponse(**mitm_dataset_bundle.to_import().model_dump())
@router.get('/mitm_dataset/{uuid}/import/zip', response_class=StreamingResponse,
- responses={200: {'content': {'application/zip': {}}}})
-def generate_uploaded_mitm_dataset_import_zip(uploaded_dataset: TrackedMitMDatasetDependency,
- include_visualizations: bool = False) -> StreamingResponse:
- assets = generate_uploaded_mitm_dataset_import(uploaded_dataset,
- include_visualizations=include_visualizations)
+ responses={200: {'content': {'application/zip': {}}}})
+def generate_tracked_mitm_dataset_import_zip(tracked_dataset: TrackedMitMDatasetDependency,
+ include_visualizations: bool = False) -> StreamingResponse:
+ assets = generate_tracked_mitm_dataset_import(tracked_dataset,
+ include_visualizations=include_visualizations)
bio = io.BytesIO()
write_superset_import_as_zip(bio, assets)
bio.seek(0)