diff --git a/app/db/models/__init__.py b/app/db/models/__init__.py
index 7b68cfcde95f352654a63c30236dec3320b304a7..47e5579d16d0dedc795ba2c61361fe78cdd069f5 100644
--- a/app/db/models/__init__.py
+++ b/app/db/models/__init__.py
@@ -1,3 +1,4 @@
-from .tracked_mitm_dataset import BaseTrackedMitMDataset, AddTrackedMitMDataset, TrackedMitMDataset
from .common import FromPydanticModelsMixin, APPLICATION_DB_SCHEMA
-from .presentation import ListTrackedMitMDataset
\ No newline at end of file
+from .tracked_mitm_dataset import BaseTrackedMitMDataset, AddTrackedMitMDataset, TrackedMitMDataset, ExternalTrackedMitMDataset
+from .mapped_sources import MappedDB, MappedDBSource, MappedDBPull
+from .presentation import ListTrackedMitMDataset
diff --git a/app/db/models/common.py b/app/db/models/common.py
index f01fa7b609c77dcf25f9598ba2a09c04fd85ff54..34ff6efa1da1de52f12c85d92bfc8d3e3b556a91 100644
--- a/app/db/models/common.py
+++ b/app/db/models/common.py
@@ -13,6 +13,6 @@ class FromPydanticModelsMixin:
def from_models(cls, *base_objs: pydantic.BaseModel, **kwargs) -> Self:
const_kwargs = {}
for base_obj in base_objs:
- const_kwargs |= base_obj.model_dump(round_trip=True)
+ const_kwargs |= base_obj.__dict__
const_kwargs |= kwargs
return cls(**const_kwargs)
diff --git a/app/db/models/mapped_sources.py b/app/db/models/mapped_sources.py
index b39ce6970ebfc8f916673745f4b78eb1018795ee..40d6107b64df3e9d6d1e0cd9596b085296d9940f 100644
--- a/app/db/models/mapped_sources.py
+++ b/app/db/models/mapped_sources.py
@@ -1,16 +1,23 @@
+from __future__ import annotations
+
import uuid
+from datetime import datetime
from uuid import UUID
import pydantic
+import sqlmodel
from mitm_tooling.definition import MITM
from mitm_tooling.extraction.sql.data_models import CompiledVirtualView
+from mitm_tooling.extraction.sql.data_models.db_meta import DBMetaInfoBase
+from mitm_tooling.extraction.sql.data_models.db_probe import DBProbeBase
from mitm_tooling.extraction.sql.mapping import ConceptMapping
-from pydantic import BaseModel
+from mitm_tooling.representation import Header
+from pydantic import BaseModel, AnyUrl
from sqlalchemy.orm import relationship
from sqlmodel import Field, SQLModel
from .common import APPLICATION_DB_SCHEMA
-from ..adapters import PydanticType
+from ..adapters import PydanticType, StrType
class MappedDB(BaseModel):
@@ -23,12 +30,56 @@ class MappedDB(BaseModel):
return [cm for cm in self.mappings if cm.mitm == self.mitm]
-class MappedMitMDatasetSources(SQLModel, table=True):
+class DBInfo(BaseModel):
+ db_meta: DBMetaInfoBase
+ db_probe: DBProbeBase
+
+
+class MappedDBPull(SQLModel, table=True):
model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
- __tablename__ = 'mapped_mitmdataset_sources'
+ __tablename__ = 'mapped_db_pulls'
+ __table_args__ = {'schema': APPLICATION_DB_SCHEMA}
+
+ id: int = Field(primary_key=True, sa_column_kwargs={'autoincrement': True})
+ tracked_mitm_dataset_id: int = Field(nullable=False,
+ foreign_key=f'{APPLICATION_DB_SCHEMA}.tracked_mitm_datasets.id')
+ mapped_db_source_id: int = Field(nullable=False, foreign_key=f'{APPLICATION_DB_SCHEMA}.mapped_db_sources.id')
+
+ time: datetime = Field(sa_type=sqlmodel.DateTime, default_factory=datetime.now)
+ instances_imported: int = Field(default=0)
+ rows_created: int = Field(default=0)
+
+ @property
+ def tracked_mitm_dataset(self) -> 'ExternalTrackedMitMDataset':
+ return relationship('ExternalTrackedMitMDataset', foreign_keys='tracked_mitm_dataset_id', lazy='joined')
+
+ @property
+ def mapped_db_source(self) -> MappedDBSource:
+ return relationship('MappedDBSource', foreign_keys='mapped_db_source_id', lazy='joined')
+
+
+class MappedDBSource(SQLModel, table=True):
+ model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+ __tablename__ = 'mapped_db_sources'
__table_args__ = {'schema': APPLICATION_DB_SCHEMA}
id: int = Field(primary_key=True, sa_column_kwargs={'autoincrement': True})
uuid: UUID = Field(default_factory=uuid.uuid4, index=True, unique=True, nullable=False)
+ sql_alchemy_uri: AnyUrl = Field(sa_type=StrType.wrap(AnyUrl))
mitm_mapping: MappedDB = Field(sa_type=PydanticType.wrap(MappedDB), repr=False)
+ mitm_header: Header = Field(sa_type=PydanticType.wrap(Header), repr=False)
+ @property
+ def pulls(self) -> list[MappedDBPull]:
+ return relationship('MappedDBPull')
+
+ # @hybridproperty
+ @property
+ def last_pulled(self) -> datetime | None:
+ from sqlalchemy import select, func
+ stmt = select(func.max(MappedDBPull.time)).where(MappedDBPull.mapped_db_source_id == self.id)
+ return stmt.scalar_subquery()
+
+ @property
+ def tracked_mitm_datasets(self) -> list['ExternalTrackedMitMDataset']:
+ return relationship('ExternalTrackedMitMDataset', back_populates='mapped_db_source')
diff --git a/app/db/models/tracked_mitm_dataset.py b/app/db/models/tracked_mitm_dataset.py
index c79c59da80e1870a8680e6f80bff1225eb7ed54d..dba89717856e6feb48cd5531d06e5e7cd7bd95bc 100644
--- a/app/db/models/tracked_mitm_dataset.py
+++ b/app/db/models/tracked_mitm_dataset.py
@@ -1,22 +1,27 @@
+from __future__ import annotations
+
import uuid
from abc import ABC
from datetime import datetime
+from typing import TYPE_CHECKING
from uuid import UUID
import pydantic
import sqlmodel
from mitm_tooling.definition import MITM
-
from mitm_tooling.representation import Header, SQLRepresentationSchema, mk_sql_rep_schema
from mitm_tooling.transformation.superset.asset_bundles import MitMDatasetIdentifierBundle, DatasourceIdentifierBundle
from mitm_tooling.transformation.superset.common import DBConnectionInfo
from mitm_tooling.transformation.superset.definitions import MitMDatasetIdentifier
from pydantic import BaseModel, AnyUrl
+from sqlalchemy.orm import relationship
from sqlmodel import SQLModel, Field
from app.db.adapters import StrType, PydanticType
from .common import FromPydanticModelsMixin, APPLICATION_DB_SCHEMA
+if TYPE_CHECKING:
+ from .mapped_sources import MappedDBSource, MappedDBPull
class BaseTrackedMitMDataset(BaseModel, ABC):
@@ -31,21 +36,26 @@ class AddTrackedMitMDataset(BaseModel):
sql_alchemy_uri: AnyUrl
mitm_header: Header
+
class TrackedMitMDataset(FromPydanticModelsMixin, AddTrackedMitMDataset, BaseTrackedMitMDataset, SQLModel, table=True):
model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
__tablename__ = 'tracked_mitm_datasets'
__table_args__ = {'schema': APPLICATION_DB_SCHEMA}
+ __mapper_args__ = {'polymorphic_on': 'type', 'polymorphic_identity': 'local'}
id: int = Field(primary_key=True, sa_column_kwargs={'autoincrement': True})
uuid: UUID = Field(default_factory=uuid.uuid4, index=True, unique=True, nullable=False)
+ type: str = Field(default='local', nullable=False)
+
dataset_name: str = Field()
schema_name: str = Field()
sql_alchemy_uri: AnyUrl = Field(sa_type=StrType.wrap(AnyUrl))
+
is_managed_locally: bool = Field(default=True)
last_edited: datetime = Field(sa_type=sqlmodel.DateTime, default_factory=datetime.now)
-
mitm_header: Header = Field(sa_type=PydanticType.wrap(Header), repr=False)
- identifier_bundle: MitMDatasetIdentifierBundle = Field(sa_type=PydanticType.wrap(MitMDatasetIdentifierBundle), repr=False)
+ identifier_bundle: MitMDatasetIdentifierBundle = Field(sa_type=PydanticType.wrap(MitMDatasetIdentifierBundle),
+ repr=False)
@property
def identifier(self) -> MitMDatasetIdentifier:
@@ -54,7 +64,8 @@ class TrackedMitMDataset(FromPydanticModelsMixin, AddTrackedMitMDataset, BaseTra
@property
def datasource_identifiers(self) -> DatasourceIdentifierBundle:
# we explicitly do not want to use the identifier_bundle itself directly, as that includes the visualization identifier map
- return DatasourceIdentifierBundle(database=self.identifier_bundle.database, ds_id_map=self.identifier_bundle.ds_id_map)
+ return DatasourceIdentifierBundle(database=self.identifier_bundle.database,
+ ds_id_map=self.identifier_bundle.ds_id_map)
@property
def db_conn_info(self) -> DBConnectionInfo:
@@ -64,7 +75,39 @@ class TrackedMitMDataset(FromPydanticModelsMixin, AddTrackedMitMDataset, BaseTra
def sql_rep_schema(self) -> SQLRepresentationSchema:
return mk_sql_rep_schema(self.mitm_header, override_schema=self.schema_name)
- #@pydantic.computed_field()
+ # @pydantic.computed_field()
@property
def mitm(self) -> MITM:
return self.mitm_header.mitm
+
+
+class ExternalTrackedMitMDataset(TrackedMitMDataset):
+ model_config = pydantic.ConfigDict(arbitrary_types_allowed=True)
+ __mapper_args__ = {'polymorphic_identity': 'mapped'}
+
+ type: str = Field(default='mapped', nullable=False)
+ sql_alchemy_uri: AnyUrl = Field(sa_type=StrType.wrap(AnyUrl))
+ is_managed_locally: bool = Field(default=False)
+ last_edited: datetime = Field(sa_type=sqlmodel.DateTime, default_factory=datetime.now)
+ mitm_header: Header = Field(sa_type=PydanticType.wrap(Header), repr=False)
+ identifier_bundle: MitMDatasetIdentifierBundle = Field(sa_type=PydanticType.wrap(MitMDatasetIdentifierBundle),
+ repr=False)
+
+ mapped_db_source_id: int = Field(foreign_key=f'{APPLICATION_DB_SCHEMA}.mapped_db_sources.id', nullable=False)
+
+ @property
+ def mapped_db_source(self) -> MappedDBSource:
+ return relationship('MappedDBSource',
+ back_populates='tracked_mitm_datasets',
+ foreign_keys='mapped_db_source_id')
+
+ @property
+ def pulls(self) -> list[MappedDBPull]:
+ return relationship('MappedDBPull', foreign_keys='tracked_mitm_dataset_id')
+
+ # @hybridproperty
+ @property
+ def last_pulled(self) -> datetime | None:
+ from sqlalchemy import select, func
+ stmt = select(func.max('MappedDBPull.time')).where('MappedDBPull.tracked_mitm_dataset_id' == self.id)
+ return stmt.scalar_subquery()
diff --git a/app/db/setup.py b/app/db/setup.py
index d077a782b5df778fcddadb49c1cb267a64b5ff94..a4d2bc2c60f95bf88c09f4b83a69cc6c535503cd 100644
--- a/app/db/setup.py
+++ b/app/db/setup.py
@@ -4,7 +4,6 @@ import sqlalchemy as sa
from mitm_tooling.utilities.python_utils import pick_from_mapping
from sqlalchemy import create_engine, inspect, Engine
from sqlalchemy.orm import Session
-from sqlmodel import SQLModel
from ..config import app_cfg
@@ -20,7 +19,9 @@ if MITM_DATABASE_URL.get_dialect().name == 'sqlite':
engine = create_engine(MITM_DATABASE_URL, execution_options=execution_options)
def init_db():
+ from .models import TrackedMitMDataset, ExternalTrackedMitMDataset, MappedDBSource, MappedDBPull
from .models import APPLICATION_DB_SCHEMA
+ from sqlmodel import SQLModel
from .utils import create_schema
logger.info(f'Setting up MITM DB @ {MITM_DATABASE_URL}')
with engine.connect() as conn:
diff --git a/app/db/utils.py b/app/db/utils.py
index a11965b665b16cb8a7e7131c96d4fcd415726199..f0c089cd8297135ac7af97c0a83ef3711f438626 100644
--- a/app/db/utils.py
+++ b/app/db/utils.py
@@ -33,7 +33,7 @@ def mk_orm_session(engine: Engine) -> Generator[ORMSession, None, None]:
def infer_tracked_mitm_dataset_schema(engine: sa.Engine, uuid: UUID) -> DBMetaInfo | None:
with mk_session(engine) as session:
- model = session.get(TrackedMitMDataset, (uuid,))
+ model = session.query(TrackedMitMDataset).filter(TrackedMitMDataset.uuid == uuid).one_or_none()
if model is not None:
meta, _ = connect_and_reflect(engine, allowed_schemas={model.schema_name})
return DBMetaInfo.from_sa_meta(meta, default_schema=model.schema_name)
diff --git a/app/dependencies/orm.py b/app/dependencies/orm.py
index c62d716f2de2d66e7ffc91a0600416000406e575..66281981e65f8c3de04e6451b94f2b1cb0797a3a 100644
--- a/app/dependencies/orm.py
+++ b/app/dependencies/orm.py
@@ -9,7 +9,7 @@ from ..db.models import TrackedMitMDataset
def get_uploaded_dataset(session: ORMSessionDependency, uuid: UUID = fastapi.Path()) -> TrackedMitMDataset:
- o = session.get(TrackedMitMDataset, (uuid,))
+ o = session.query(TrackedMitMDataset).filter(TrackedMitMDataset.uuid == uuid).one_or_none()
if o is None:
raise HTTPException(status_code=404, detail='Referenced MitM Dataset does not exist.')
return o
diff --git a/app/routes/definitions/generate.py b/app/routes/definitions/generate.py
index 1533ddd322b6842ef6d4027bf442c7aed20d545e..1f2992eabb6fe8a2753d55b6bbc877dc9d187b95 100644
--- a/app/routes/definitions/generate.py
+++ b/app/routes/definitions/generate.py
@@ -1,5 +1,5 @@
-from datetime import datetime
from collections.abc import Sequence
+from datetime import datetime
from mitm_tooling.definition import MITM
from mitm_tooling.representation import Header
@@ -8,7 +8,7 @@ from mitm_tooling.transformation.superset import mk_superset_mitm_dataset_bundle
from mitm_tooling.transformation.superset.asset_bundles import MitMDatasetIdentifierBundle, SupersetMitMDatasetBundle, \
SupersetDatasourceBundle, SupersetVisualizationBundle
from mitm_tooling.transformation.superset.common import DBConnectionInfo
-from mitm_tooling.transformation.superset.definitions import SupersetAssetsImport, SupersetMitMDatasetImport, \
+from mitm_tooling.transformation.superset.definitions import SupersetMitMDatasetImport, \
MetadataType
from app.db.models import TrackedMitMDataset
@@ -120,7 +120,8 @@ def track_visualizations(orm_session: ORMSession,
def exec_viz_import_request(orm_session: ORMSession,
tracked_dataset: TrackedMitMDataset,
- request: GenerateVisualizationsRequest, as_assets: bool = False) -> SupersetMitMDatasetImport:
+ request: GenerateVisualizationsRequest) -> SupersetMitMDatasetImport:
mitm_dataset_bundle = exec_viz_request(orm_session, tracked_dataset, request)
- importable = mitm_dataset_bundle.to_import(metadata_type=(MetadataType.Asset if as_assets else MetadataType.MitMDataset))
+ importable = mitm_dataset_bundle.to_import(metadata_type=(
+ request.override_metadata_type if request.override_metadata_type else MetadataType.MitMDataset))
return importable
diff --git a/app/routes/definitions/requests.py b/app/routes/definitions/requests.py
index 060c12e73c9f43578a2e0bd5aae336f42a8e1352..f673d5d0e7b2f2fd7fef47d37f909fdacdcbc41b 100644
--- a/app/routes/definitions/requests.py
+++ b/app/routes/definitions/requests.py
@@ -3,7 +3,7 @@ from mitm_tooling.representation import Header
from mitm_tooling.transformation.superset import VisualizationType, MAEDVisualizationType
from mitm_tooling.transformation.superset.asset_bundles import MitMDatasetIdentifierBundle
from mitm_tooling.transformation.superset.common import DBConnectionInfo, MitMDatasetInfo
-from mitm_tooling.transformation.superset.definitions import StrUUID
+from mitm_tooling.transformation.superset.definitions import StrUUID, MetadataType
from mitm_tooling.transformation.superset.definitions.mitm_dataset import MitMDatasetIdentifier
@@ -18,3 +18,4 @@ class GenerateVisualizationsRequest(pydantic.BaseModel):
visualization_types: list[VisualizationType] = [MAEDVisualizationType.Baseline]
reuse_existing_identifiers: bool = True
track_identifiers: bool = False
+ override_metadata_type: MetadataType | None = None
diff --git a/app/routes/definitions/router.py b/app/routes/definitions/router.py
index 02938ea32750e8540353a9d154a2dee6caf74822..193c098e0357cc50cfdf0e42a1b1b43b6b25f3fb 100644
--- a/app/routes/definitions/router.py
+++ b/app/routes/definitions/router.py
@@ -61,7 +61,7 @@ def generate_tracked_mitm_dataset_import(tracked_dataset: TrackedMitMDatasetDepe
importable = exec_tracked_asset_import_request(tracked_dataset,
include_visualizations,
override_metadata_type=override_metadata_type)
- return MitMDatasetImportResponse.from_models(importable.model_dump())
+ return MitMDatasetImportResponse.from_models(importable)
@router.get('/mitm_dataset/{uuid}/import/zip', response_class=StreamingResponse,
@@ -84,25 +84,23 @@ def generate_visualizations_for_tracked_dataset(orm_session: ORMSessionDependenc
tracked_dataset: TrackedMitMDatasetDependency,
request: GenerateVisualizationsRequest) -> MitMDatasetBundleResponse:
mitm_dataset_bundle = exec_viz_request(orm_session, tracked_dataset, request)
- return MitMDatasetBundleResponse.from_models(mitm_dataset_bundle.model_dump())
+ return MitMDatasetBundleResponse.from_models(mitm_dataset_bundle)
@router.post('/mitm_dataset/viz/{uuid}/import')
def generate_visualizations_import_for_tracked_dataset(orm_session: ORMSessionDependency,
tracked_dataset: TrackedMitMDatasetDependency,
- request: GenerateVisualizationsRequest,
- as_assets: bool = False) -> VisualizationImportResponse:
- importable = exec_viz_import_request(orm_session, tracked_dataset, request, as_assets=as_assets)
- return VisualizationImportResponse.from_models(importable.model_dump())
+ request: GenerateVisualizationsRequest) -> VisualizationImportResponse:
+ importable = exec_viz_import_request(orm_session, tracked_dataset, request)
+ return VisualizationImportResponse.from_models(importable)
@router.post('/mitm_dataset/viz/{uuid}/import/zip', response_class=StreamingResponse,
responses={200: {'content': {'application/zip': {}}}})
def generate_visualizations_import_zip_for_tracked_dataset(orm_session: ORMSessionDependency,
tracked_dataset: TrackedMitMDatasetDependency,
- request: GenerateVisualizationsRequest,
- as_assets: bool = False) -> StreamingResponse:
- importable = exec_viz_import_request(orm_session, tracked_dataset, request, as_assets=as_assets)
+ request: GenerateVisualizationsRequest) -> StreamingResponse:
+ importable = exec_viz_import_request(orm_session, tracked_dataset, request)
bio = io.BytesIO()
write_superset_import_as_zip(bio, importable)
bio.seek(0)
diff --git a/app/routes/mitm_dataset/export.py b/app/routes/mitm_dataset/export.py
index c68711762353bf55528cc3bc3659ebf4e1b3a520..83185bd4dd29e07c7f2a58d2b6e8a2ea079efe2b 100644
--- a/app/routes/mitm_dataset/export.py
+++ b/app/routes/mitm_dataset/export.py
@@ -9,6 +9,7 @@ from mitm_tooling.transformation.sql.into_mappings import sql_rep_into_mappings
from mitm_tooling.utilities.sql_utils import create_sa_engine
from app.db.models import TrackedMitMDataset
+from app.routes.mitm_dataset.mapped_db import mk_exportable
logger = logging.getLogger(__name__)
diff --git a/app/routes/mitm_dataset/register_external.py b/app/routes/mitm_dataset/register_external.py
index a50e08ac367c4cb01dbb8a72f4569b757e6fc262..8ea786f965631addbe054ab5b7e1b06da4c89817 100644
--- a/app/routes/mitm_dataset/register_external.py
+++ b/app/routes/mitm_dataset/register_external.py
@@ -1,3 +1,5 @@
+from datetime import datetime
+
import sqlalchemy as sa
from mitm_tooling.definition import MITM
from mitm_tooling.extraction.sql.data_models import VirtualDB, DBMetaInfo, SourceDBType, VirtualView, \
@@ -12,19 +14,51 @@ from mitm_tooling.utilities.sql_utils import create_sa_engine
from pydantic import AnyUrl, BaseModel
from sqlalchemy.orm import Session
+from app.db.models.mapped_sources import MappedDBSource, MappedDBPull
+from app.db.models.tracked_mitm_dataset import ExternalTrackedMitMDataset
from app.dependencies.db import ORMSessionDependency
+from app.routes.mitm_dataset.append import append_exportable
+from app.routes.mitm_dataset.mapped_db import mk_exportable
from app.routes.mitm_dataset.requests import RegisterExternalMitMDatasetRequest
-from app.routes.mitm_dataset.responses import RegisterMitMResponse
-
+from app.routes.mitm_dataset.responses import RegisterExternalMitMResponse
+from app.routes.mitm_dataset.upload import upload_exportable
def register_external_mitm_dataset(
session: ORMSessionDependency,
request: RegisterExternalMitMDatasetRequest,
-) -> RegisterMitMResponse:
+) -> RegisterExternalMitMResponse:
+
+
+ remote_engine = create_sa_engine(request.sql_alchemy_uri)
+ exportable = mk_exportable(remote_engine, request.mapped_db)
+ # header = exportable.generate_header(remote_engine)
+
+ add_tracked_mitm_dataset = upload_exportable(request.sql_alchemy_uri, exportable, request.dataset_name, skip_instances=True)
+
+ mapped_db_source = MappedDBSource(sql_alchemy_uri=request.sql_alchemy_uri, mitm_mapping=request.mapped_db, mitm_header=add_tracked_mitm_dataset.mitm_header)
+ session.add(mapped_db_source)
+
+ external_tracked_mitm_dataset = ExternalTrackedMitMDataset.from_models(add_tracked_mitm_dataset, mapped_db_source=mapped_db_source)
+
+ session.add(external_tracked_mitm_dataset)
+
+ session.commit()
+ session.refresh(external_tracked_mitm_dataset)
+ session.refresh(mapped_db_source)
+ return RegisterExternalMitMResponse(tracked_mitm_dataset=external_tracked_mitm_dataset, mapped_db_source=mapped_db_source)
+
+def pull(session: ORMSessionDependency, external_tracked_mitm_dataset: ExternalTrackedMitMDataset):
+ db_source = external_tracked_mitm_dataset.mapped_db_source
+
+ remote_engine = create_sa_engine(db_source.sql_alchemy_uri)
+
+ exportable = mk_exportable(remote_engine, db_source.mitm_mapping)
- sql_rep_schema = mk_sql_rep_schema(header, view_generators=None)
+ ii, ir = append_exportable(db_source.sql_alchemy_uri, exportable, external_tracked_mitm_dataset)
+ external_tracked_mitm_dataset.last_edited = datetime.now()
- # TODO connect mapped_vvs to sql_rep_schema
+ pull = MappedDBPull(tracked_mitm_dataset_id=external_tracked_mitm_dataset.id, mapped_db_source_id=db_source.id, instances_imported=ii, rows_created=ir)
- raise NotImplementedError()
+ session.add(pull)
+ session.flush()
diff --git a/app/routes/mitm_dataset/requests.py b/app/routes/mitm_dataset/requests.py
index 877dfa73fae66727398c3014c198a7245114e2a8..5ad67111e48f5f7484aba4f799ade1c8d9f5a6df 100644
--- a/app/routes/mitm_dataset/requests.py
+++ b/app/routes/mitm_dataset/requests.py
@@ -1,16 +1,17 @@
+from uuid import UUID
+
import pydantic
-from mitm_tooling.definition import MITM
-from mitm_tooling.extraction.sql.data_models import CompiledVirtualView
-from mitm_tooling.extraction.sql.mapping import ConceptMapping
from mitm_tooling.representation import Header
from pydantic import AnyUrl
from app.db.models import AddTrackedMitMDataset
+from app.db.models.mapped_sources import MappedDB
class AddTrackedMitMDatasetRequest(AddTrackedMitMDataset):
pass
+
class EditTrackedMitMDatasetRequest(pydantic.BaseModel):
dataset_name: str
schema_name: str
@@ -18,9 +19,9 @@ class EditTrackedMitMDatasetRequest(pydantic.BaseModel):
mitm_header: Header
is_managed_locally: bool
+
class RegisterExternalMitMDatasetRequest(pydantic.BaseModel):
+ uuid: UUID | None = None
dataset_name: str
sql_alchemy_uri: AnyUrl
- mitm: MITM
- cvvs: list[CompiledVirtualView]
- mappings: list[ConceptMapping]
+ mapped_db: MappedDB
diff --git a/app/routes/mitm_dataset/responses.py b/app/routes/mitm_dataset/responses.py
index 09daed4a0974d0f8a1152da5e12414ade02e439e..3813a1c62dc2239dc22dc71944002e4d13cb0cb2 100644
--- a/app/routes/mitm_dataset/responses.py
+++ b/app/routes/mitm_dataset/responses.py
@@ -7,6 +7,7 @@ from mitm_tooling.definition import MITM
from pydantic import Field
from app.db.models import TrackedMitMDataset, BaseTrackedMitMDataset, ListTrackedMitMDataset
+from app.db.models.mapped_sources import MappedDBSource
class TrackMitMResponse(pydantic.BaseModel):
@@ -19,8 +20,8 @@ class UploadMitMResponse(TrackMitMResponse):
pass
-class RegisterMitMResponse(TrackMitMResponse):
- pass
+class RegisterExternalMitMResponse(TrackMitMResponse):
+ mapped_db_source: MappedDBSource | None = None
MitMsListResponse = pydantic.TypeAdapter(list[ListTrackedMitMDataset])
diff --git a/app/routes/mitm_dataset/router.py b/app/routes/mitm_dataset/router.py
index 6b17a20c2ca258c7d4c8c2014acccdbfa52b5708..fdbb2794f0fee30ebf2f2c2643ae2ee1c12f1374 100644
--- a/app/routes/mitm_dataset/router.py
+++ b/app/routes/mitm_dataset/router.py
@@ -15,7 +15,7 @@ from app.dependencies.db import DBEngineDependency, ORMSessionDependency
from app.dependencies.orm import TrackedMitMDatasetDependency
from .export import export_via_mapping
from .requests import AddTrackedMitMDatasetRequest, RegisterExternalMitMDatasetRequest, EditTrackedMitMDatasetRequest
-from .responses import UploadMitMResponse, RegisterMitMResponse
+from .responses import UploadMitMResponse, RegisterExternalMitMResponse
from .upload import upload_mitm_file
from ...db.utils import mk_session
@@ -44,7 +44,7 @@ def upload_mitm_dataset(
@router.post('/register')
def register_external_mitm_dataset(session: ORMSessionDependency,
- request: RegisterExternalMitMDatasetRequest) -> RegisterMitMResponse:
+ request: RegisterExternalMitMDatasetRequest) -> RegisterExternalMitMResponse:
from .register_external import register_external_mitm_dataset
return register_external_mitm_dataset(session, request)
@@ -53,7 +53,7 @@ def register_external_mitm_dataset(session: ORMSessionDependency,
def post_mitm_dataset(session: ORMSessionDependency,
new_mitm_dataset: AddTrackedMitMDatasetRequest) -> TrackedMitMDataset:
try:
- new = TrackedMitMDataset.model_validate(**new_mitm_dataset.model_dump(mode='python', round_trip=True))
+ new = TrackedMitMDataset.from_models(new_mitm_dataset)
session.add(new)
session.commit()
session.refresh(new)
@@ -103,11 +103,11 @@ def export_mitm_dataset(engine: DBEngineDependency,
tracked_dataset: TrackedMitMDatasetDependency,
use_streaming: bool = False) -> StreamingResponse:
remote_engine, exportable = export_via_mapping(tracked_dataset)
- with mk_session(remote_engine) as session:
+ with remote_engine.begin() as conn:
if use_streaming:
- ze = exportable.export_as_stream(session)
+ ze = exportable.export_as_stream(conn)
data = ze.iter_bytes()
else:
- ze = exportable.export_to_memory(session)
+ ze = exportable.export_to_memory(conn)
data = ze.to_buffer()
return StreamingResponse(data, media_type='application/zip', headers={'Content-Disposition': 'attachment; filename=export.zip'})
diff --git a/app/routes/mitm_dataset/upload.py b/app/routes/mitm_dataset/upload.py
index 5ae32c70c1bc5777fff1fc4253a9fccfa0135ffa..b394d93dfe2573af48e01b200fc3ae3411cc5443 100644
--- a/app/routes/mitm_dataset/upload.py
+++ b/app/routes/mitm_dataset/upload.py
@@ -6,10 +6,10 @@ import sqlalchemy as sa
from mitm_tooling.definition import MITM
from mitm_tooling.extraction.sql.mapping import Exportable
from mitm_tooling.io import read_zip
-from mitm_tooling.representation import mk_sql_rep_schema, MITMData, insert_mitm_data, Header, SQLRepresentationSchema
+from mitm_tooling.representation import mk_sql_rep_schema, MITMData, Header, SQLRepresentationSchema
from mitm_tooling.representation.sql_representation import insert_db_schema, insert_mitm_data_instances
-from mitm_tooling.transformation.sql.from_exportable import insert_exportable, insert_exportable_instances
-from mitm_tooling.utilities.identifiers import name_plus_uuid, mk_uuid
+from mitm_tooling.transformation.sql.from_exportable import insert_exportable_instances
+from mitm_tooling.utilities.identifiers import name_plus_uuid, mk_uuid, mk_short_uuid_str
from mitm_tooling.utilities.io_utils import DataSource
from mitm_tooling.utilities.sql_utils import sa_url_into_any_url, create_sa_engine
from pydantic import AnyUrl
@@ -22,6 +22,16 @@ from app.dependencies.db import get_engine
logger = logging.getLogger(__name__)
+def _skip_instances(
+ conn: sa.Connection,
+ sql_rep_schema: SQLRepresentationSchema,
+) -> tuple[int, int]:
+ """
+ Skip the instances insertion.
+ """
+ return 0, 0
+
+
def upload_mitm_file(mitm: MITM,
mitm_zip: DataSource,
dataset_name: str,
@@ -34,9 +44,13 @@ def upload_mitm_file(mitm: MITM,
def upload_mitm_data(mitm_data: MITMData,
dataset_name: str,
uuid: UUID | None = None,
- engine: Engine = None) -> AddTrackedMitMDataset:
- header_creator = lambda : mitm_data.header
- instances_inserter = lambda conn, sql_rep_schema: insert_mitm_data_instances(conn, sql_rep_schema, mitm_data)
+ engine: Engine = None,
+ skip_instances: bool = False) -> AddTrackedMitMDataset:
+ header_creator = lambda: mitm_data.header
+ if skip_instances:
+ instances_inserter = _skip_instances
+ else:
+ instances_inserter = lambda conn, sql_rep_schema: insert_mitm_data_instances(conn, sql_rep_schema, mitm_data)
return upload_data(header_creator, instances_inserter, dataset_name, uuid, engine)
@@ -45,13 +59,17 @@ def upload_exportable(source: AnyUrl,
exportable: Exportable,
dataset_name: str,
uuid: UUID | None = None,
- engine: Engine = None) -> AddTrackedMitMDataset:
+ engine: Engine = None,
+ skip_instances: bool = False) -> AddTrackedMitMDataset:
source_engine = create_sa_engine(source)
header_creator = lambda: exportable.generate_header(source_engine)
- def instances_inserter(conn, sql_rep_schema):
- return insert_exportable_instances(source_engine, exportable, conn, sql_rep_schema, stream_data=False)
+ if skip_instances:
+ instances_inserter = _skip_instances
+ else:
+ def instances_inserter(conn, sql_rep_schema):
+ return insert_exportable_instances(source_engine, exportable, conn, sql_rep_schema, stream_data=False)
return upload_data(header_creator, instances_inserter, dataset_name, uuid, engine)
@@ -64,7 +82,7 @@ def upload_data(header_creator: Callable[[], Header],
engine = engine if engine is not None else get_engine()
sql_alchemy_uri = sa_url_into_any_url(engine.url)
uuid = uuid or mk_uuid()
- unique_schema_name = name_plus_uuid(dataset_name, uuid, sep='_')
+ unique_schema_name = mk_short_uuid_str(uuid) # name_plus_uuid(dataset_name, uuid, sep='_')
logger.info(f'Uploading MitM Data with uuid {uuid} into target schema {unique_schema_name} on connected DB {engine.url}.')
diff --git a/docker-compose.yaml b/docker-compose.yaml
index 2e6b30fdbe08b1f0e5f7f31fff697169f5e75216..4a9570449ab47bdbe108980c56f7b95071255299 100644
--- a/docker-compose.yaml
+++ b/docker-compose.yaml
@@ -22,7 +22,7 @@ services:
container_name: test_mitm_db
restart: unless-stopped
ports:
- - "5432:5432"
+ - "25432:5432"
env_file: docker/.env
volumes:
- mitm_db_home:/var/lib/postgresql/data
diff --git a/test/http-client.env.json b/test/http-client.env.json
index 5f385fd7c7327f7b2bc979fed5ad5bfcdeeee70a..a70cd6df712c2411f974707758f9591cf194f8d4 100644
--- a/test/http-client.env.json
+++ b/test/http-client.env.json
@@ -9,7 +9,7 @@
},
"superset": {
"port": "8180",
- "uuid": "5b9bc5c0-7faf-4807-a098-573558a0e2a6"
+ "uuid": "6ea60ebf-b3ec-4b64-988d-8dc4e2f48eca"
},
"kubernetes": {
"port": "8080",
diff --git a/test/upload.http b/test/upload.http
index 706b92cb3cf2a5daeb1dcb9927eae5c2b5ce2864..e62824c1f60e47004ecd0b3aca51934754105875 100644
--- a/test/upload.http
+++ b/test/upload.http
@@ -3,7 +3,7 @@
# @name Upload MAED dataset
# @timeout 180
# @connection-timeout 180
-POST http://localhost:{{port}}/mitm_dataset/upload?dataset_name=myname_0&mitm=MAED
+POST http://localhost:{{port}}/mitm_dataset/upload?dataset_name=myname_x&mitm=MAED
Accept: application/json
Content-Type: multipart/form-data; boundary=WebAppBoundary