From 7ea67895f2748018e885cdfff140b048dbe92c45 Mon Sep 17 00:00:00 2001
From: Leah Tacke genannt Unterberg <leah.tgu@pads.rwth-aachen.de>
Date: Thu, 10 Apr 2025 16:58:33 +0200
Subject: [PATCH] fixed some things in the api and tasks
---
justfile | 2 +-
requirements/base.txt | 1 +
requirements/development.txt | 12 ++--
.../hooks}/asyncExternalService.ts | 2 +-
.../UploadMitMDatasetModal/index.tsx | 2 +-
.../src/pages/MitMDatasetList/index.tsx | 57 ++++++++++++++++---
superset/commands/mitm/mitm_dataset/delete.py | 2 +-
superset/commands/mitm/mitm_service/delete.py | 4 +-
superset/config.py | 2 +
superset/customization/mitm_datasets/api.py | 48 +++++++++++++---
superset/tasks/mitm/__init__.py | 4 ++
superset/tasks/mitm/call_external_service.py | 3 +-
superset/tasks/mitm/download_mitm_dataset.py | 8 +--
superset/tasks/mitm/drop_mitm_dataset.py | 12 ++--
superset/tasks/mitm/upload_mitm_dataset.py | 8 +--
tests/requests/mitm_dataset_api.http | 9 +++
16 files changed, 135 insertions(+), 41 deletions(-)
rename superset-frontend/src/features/{mitmDatasets => externalServices/hooks}/asyncExternalService.ts (98%)
diff --git a/justfile b/justfile
index 09a32e10e4..fd99d8d871 100644
--- a/justfile
+++ b/justfile
@@ -1,7 +1,7 @@
set windows-shell := ["pwsh", "-c"]
up:
- docker compose up
+ docker compose up --build
pyvenv:
./.venv/Scripts/python.exe -m pip install -r requirements/development.txt
diff --git a/requirements/base.txt b/requirements/base.txt
index 24a8442a62..5d921416e3 100644
--- a/requirements/base.txt
+++ b/requirements/base.txt
@@ -151,6 +151,7 @@ greenlet==3.1.1
# via
# apache-superset (pyproject.toml)
# shillelagh
+ # sqlalchemy
gunicorn==23.0.0
# via apache-superset (pyproject.toml)
h11==0.14.0
diff --git a/requirements/development.txt b/requirements/development.txt
index ddf315bfa5..0802dfa87d 100644
--- a/requirements/development.txt
+++ b/requirements/development.txt
@@ -110,7 +110,7 @@ click-didyoumean==0.3.1
# via
# -c requirements/base.txt
# celery
-click-option-group==0.5.6
+click-option-group==0.5.7
# via
# -c requirements/base.txt
# apache-superset
@@ -200,7 +200,7 @@ flask==2.3.3
# flask-sqlalchemy
# flask-testing
# flask-wtf
-flask-appbuilder==4.6.0
+flask-appbuilder==4.6.1
# via
# -c requirements/base.txt
# apache-superset
@@ -340,7 +340,7 @@ holidays==0.25
# -c requirements/base.txt
# apache-superset
# prophet
-humanize==4.12.1
+humanize==4.12.2
# via
# -c requirements/base.txt
# apache-superset
@@ -439,7 +439,7 @@ marshmallow==3.26.1
# -c requirements/base.txt
# flask-appbuilder
# marshmallow-sqlalchemy
-marshmallow-sqlalchemy==1.3.0
+marshmallow-sqlalchemy==1.4.0
# via
# -c requirements/base.txt
# flask-appbuilder
@@ -554,7 +554,7 @@ pillow==10.4.0
# via
# apache-superset
# matplotlib
-platformdirs==4.3.6
+platformdirs==4.3.7
# via
# -c requirements/base.txt
# black
@@ -653,7 +653,7 @@ pyopenssl==25.0.0
# via
# -c requirements/base.txt
# shillelagh
-pyparsing==3.2.1
+pyparsing==3.2.2
# via
# -c requirements/base.txt
# apache-superset
diff --git a/superset-frontend/src/features/mitmDatasets/asyncExternalService.ts b/superset-frontend/src/features/externalServices/hooks/asyncExternalService.ts
similarity index 98%
rename from superset-frontend/src/features/mitmDatasets/asyncExternalService.ts
rename to superset-frontend/src/features/externalServices/hooks/asyncExternalService.ts
index f32b5ab31b..983adbf04d 100644
--- a/superset-frontend/src/features/mitmDatasets/asyncExternalService.ts
+++ b/superset-frontend/src/features/externalServices/hooks/asyncExternalService.ts
@@ -8,7 +8,7 @@ import {
SupersetClient,
SupersetError,
} from '@superset-ui/core';
-import {useAsyncEventHandling} from "../../middleware/asyncEvent";
+import {useAsyncEventHandling} from "../../../middleware/asyncEvent";
const extServiceCallApiEndpoint = '/api/v1/ext_service/call';
diff --git a/superset-frontend/src/features/mitmDatasets/UploadMitMDatasetModal/index.tsx b/superset-frontend/src/features/mitmDatasets/UploadMitMDatasetModal/index.tsx
index 7ef54e1c88..737a6f45fc 100644
--- a/superset-frontend/src/features/mitmDatasets/UploadMitMDatasetModal/index.tsx
+++ b/superset-frontend/src/features/mitmDatasets/UploadMitMDatasetModal/index.tsx
@@ -9,7 +9,7 @@ import {
formStyles,
} from '../../databases/UploadDataModel/styles';
import { Form, Input, Select, Upload } from 'antd-v5';
-import { AsyncActionResult, useExternalService } from '../asyncExternalService';
+import { AsyncActionResult, useExternalService } from '../../externalServices/hooks/asyncExternalService';
import { InboxOutlined } from '@ant-design/icons';
interface UploadMitMDatasetModalProps {
diff --git a/superset-frontend/src/pages/MitMDatasetList/index.tsx b/superset-frontend/src/pages/MitMDatasetList/index.tsx
index 75210c86c0..e78607f420 100644
--- a/superset-frontend/src/pages/MitMDatasetList/index.tsx
+++ b/superset-frontend/src/pages/MitMDatasetList/index.tsx
@@ -50,6 +50,7 @@ import {
import { ModifiedInfo } from 'src/components/AuditInfo';
import MitMDatasetModal from '../../features/mitmDatasets/MitMDatasetModal';
import UploadMitMDatasetModal from '../../features/mitmDatasets/UploadMitMDatasetModal';
+import CopyToClipboard from '../../components/CopyToClipboard';
const Actions = styled.div`
color: ${({ theme }) => theme.colors.grayscale.base};
@@ -121,6 +122,18 @@ const MitMDatasetList: FunctionComponent<MitMDatasetListProps> = ({
const [uploadModalOpen, setUploadModalOpen] = useState<boolean>(false);
const [preparingExport, setPreparingExport] = useState<boolean>(false);
+ const [passwordFields, setPasswordFields] = useState<string[]>([]);
+ const [sshTunnelPasswordFields, setSSHTunnelPasswordFields] = useState<
+ string[]
+ >([]);
+ const [sshTunnelPrivateKeyFields, setSSHTunnelPrivateKeyFields] = useState<
+ string[]
+ >([]);
+ const [
+ sshTunnelPrivateKeyPasswordFields,
+ setSSHTunnelPrivateKeyPasswordFields,
+ ] = useState<string[]>([]);
+
const openDatasetImportModal = () => {
showImportModal(true);
};
@@ -235,6 +248,21 @@ const MitMDatasetList: FunctionComponent<MitMDatasetListProps> = ({
accessor: 'changed_on_delta_humanized',
size: 'xl',
},
+ {
+ Cell: ({
+ row: {
+ original: { mitm_header },
+ },
+ }: any) => (
+ <CopyToClipboard
+ text={JSON.stringify(mitm_header, null, 2)}
+ wrapped={false}
+ />
+ ),
+ accessor: 'mitm_header',
+ hidden: true,
+ disableSortBy: true,
+ },
{
Cell: ({ row: { original } }: any) => {
// Verify owner or isAdmin
@@ -474,12 +502,13 @@ const MitMDatasetList: FunctionComponent<MitMDatasetListProps> = ({
setDatasetCurrentlyEditing(null);
};
- const handleDatasetDelete = ({
- id,
- dataset_name: datasetName,
- }: MitMDataset) => {
+ const handleDatasetDelete = (
+ { id, dataset_name: datasetName }: MitMDataset,
+ deleteRelated: boolean = false,
+ ) => {
+ const endpoint = deleteRelated ? '/delete_with_related' : '';
SupersetClient.delete({
- endpoint: `/api/v1/mitm_dataset/${id}`,
+ endpoint: `/api/v1/mitm_dataset${endpoint}/${id}`,
}).then(
() => {
refreshData();
@@ -494,9 +523,13 @@ const MitMDatasetList: FunctionComponent<MitMDatasetListProps> = ({
);
};
- const handleBulkDatasetDelete = (datasetsToDelete: MitMDataset[]) => {
+ const handleBulkDatasetDelete = (
+ datasetsToDelete: MitMDataset[],
+ deleteRelated: boolean = false,
+ ) => {
+ const endpoint = deleteRelated ? '/delete_with_related' : '';
SupersetClient.delete({
- endpoint: `/api/v1/mitm_dataset/?q=${rison.encode(
+ endpoint: `/api/v1/mitm_dataset${endpoint}/?q=${rison.encode(
datasetsToDelete.map(({ id }) => id),
)}`,
}).then(
@@ -607,6 +640,16 @@ const MitMDatasetList: FunctionComponent<MitMDatasetListProps> = ({
onModelImport={handleDatasetImport}
show={importingDataset}
onHide={closeDatasetImportModal}
+ passwordFields={passwordFields}
+ setPasswordFields={setPasswordFields}
+ sshTunnelPasswordFields={sshTunnelPasswordFields}
+ setSSHTunnelPasswordFields={setSSHTunnelPasswordFields}
+ sshTunnelPrivateKeyFields={sshTunnelPrivateKeyFields}
+ setSSHTunnelPrivateKeyFields={setSSHTunnelPrivateKeyFields}
+ sshTunnelPrivateKeyPasswordFields={sshTunnelPrivateKeyPasswordFields}
+ setSSHTunnelPrivateKeyPasswordFields={
+ setSSHTunnelPrivateKeyPasswordFields
+ }
/>
{preparingExport && <Loading />}
diff --git a/superset/commands/mitm/mitm_dataset/delete.py b/superset/commands/mitm/mitm_dataset/delete.py
index e2975a5fc8..80c01f3941 100644
--- a/superset/commands/mitm/mitm_dataset/delete.py
+++ b/superset/commands/mitm/mitm_dataset/delete.py
@@ -13,7 +13,7 @@ from ...chart.delete import DeleteChartCommand
from ...database.delete import DeleteDatabaseCommand
from ...dashboard.delete import DeleteDashboardCommand
-class DeleteMitMDatasetCommand(BaseCommand):
+class DeleteMitMDatasetsCommand(BaseCommand):
def __init__(self, model_ids: list[int], delete_related: bool = False):
self._model_ids = model_ids
self._models: list[MitMDataset] | None = None
diff --git a/superset/commands/mitm/mitm_service/delete.py b/superset/commands/mitm/mitm_service/delete.py
index cfd78b2a3e..6033b10587 100644
--- a/superset/commands/mitm/mitm_service/delete.py
+++ b/superset/commands/mitm/mitm_service/delete.py
@@ -17,8 +17,8 @@ class DeleteUploadedMitMDatasetCommand(MitMDatasetBaseCommand):
id = self._model.id
uuid = str(self._model.uuid)
- from superset.commands.mitm.mitm_dataset.delete import DeleteMitMDatasetCommand
- DeleteMitMDatasetCommand([self._model.id], delete_related=True).run()
+ from superset.commands.mitm.mitm_dataset.delete import DeleteMitMDatasetsCommand
+ DeleteMitMDatasetsCommand([self._model.id], delete_related=True).run()
delete_upload_request = ForwardableRequest(method='DELETE',
headers=[],
diff --git a/superset/config.py b/superset/config.py
index ced2f56c89..f9f874ed68 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -1060,6 +1060,8 @@ class CeleryConfig: # pylint: disable=too-few-public-methods
# },
}
+if FEATURE_FLAGS.get('MITM_SUPPORT'):
+ CeleryConfig.imports = CeleryConfig.imports + ("superset.tasks.mitm",)
CELERY_CONFIG: type[CeleryConfig] = CeleryConfig
diff --git a/superset/customization/mitm_datasets/api.py b/superset/customization/mitm_datasets/api.py
index 9c0aa9af8e..3d0bce9f1b 100644
--- a/superset/customization/mitm_datasets/api.py
+++ b/superset/customization/mitm_datasets/api.py
@@ -30,7 +30,7 @@ from ...commands.importers.exceptions import NoValidFilesFoundError
from ...commands.importers.v1.utils import get_contents_from_bundle
from ...commands.mitm.exceptions import *
from ...commands.mitm.mitm_dataset.create import CreateMitMDatasetCommand
-from ...commands.mitm.mitm_dataset.delete import DeleteMitMDatasetCommand
+from ...commands.mitm.mitm_dataset.delete import DeleteMitMDatasetsCommand
from ...commands.mitm.mitm_dataset.export import ExportMitMDatasetsCommand
from ...commands.mitm.mitm_dataset.importers.v1 import ImportMitMDatasetsCommand
from ...commands.mitm.mitm_dataset.update import UpdateMitMDatasetCommand
@@ -78,6 +78,7 @@ class MitMDatasetRestApi(BaseSupersetModelRestApi):
RouteMethod.RELATED,
'bulk_delete'
'delete_with_related'
+ 'bulk_delete_with_related'
'upload'
}
@@ -85,6 +86,7 @@ class MitMDatasetRestApi(BaseSupersetModelRestApi):
class_permission_name = 'MitMDataset'
method_permission_name = MODEL_API_RW_METHOD_PERMISSION_MAP | {
'delete_with_related': 'write',
+ 'bulk_delete_with_related': 'write',
'upload': 'write',
'export': 'read'
}
@@ -118,7 +120,7 @@ class MitMDatasetRestApi(BaseSupersetModelRestApi):
'dashboards.uuid',
'dashboards.dashboard_title',
]
- show_select_columns = show_columns + ['slices.datasource_id', 'slices.datasource_type']
+ show_select_columns = show_columns + ['slices.datasource_id', 'slices.datasource_type'] # this is necessary due to eager query manipuation in the slice model
list_columns = [
'id',
'dataset_name',
@@ -130,6 +132,9 @@ class MitMDatasetRestApi(BaseSupersetModelRestApi):
'changed_by.last_name',
'created_by.first_name',
'created_by.last_name',
+ 'owners.id',
+ 'owners.first_name',
+ 'owners.last_name',
'database.database_name',
'tables.id',
'tables.table_name',
@@ -162,7 +167,7 @@ class MitMDatasetRestApi(BaseSupersetModelRestApi):
'changed_by',
]
- allowed_rel_fields = {'database', 'creator', 'tables', 'slices', 'dashboards',
+ allowed_rel_fields = {'database', 'creator', 'owners', 'tables', 'slices', 'dashboards',
'created_by', 'changed_by'}
base_related_field_filters = {
'owners': [['id', BaseFilterRelatedUsers, lambda: []]],
@@ -335,7 +340,7 @@ class MitMDatasetRestApi(BaseSupersetModelRestApi):
)
def delete(self, pk: int) -> Response:
try:
- DeleteMitMDatasetCommand([pk]).run()
+ DeleteMitMDatasetsCommand([pk]).run()
return self.response(200, message='OK')
except MitMDatasetNotFoundError:
return self.response_404()
@@ -361,7 +366,7 @@ class MitMDatasetRestApi(BaseSupersetModelRestApi):
)
def delete_with_related(self, pk: int) -> Response:
try:
- DeleteMitMDatasetCommand([pk], delete_related=True).run()
+ DeleteMitMDatasetsCommand([pk], delete_related=True).run()
return self.response(200, message='OK')
except MitMDatasetNotFoundError:
return self.response_404()
@@ -369,7 +374,7 @@ class MitMDatasetRestApi(BaseSupersetModelRestApi):
return self.response_403()
except MitMDatasetDeleteFailedError as ex:
logger.error(
- 'Error deleting model %s: %s',
+ 'Error deleting model with its related objects %s: %s',
self.__class__.__name__,
str(ex),
exc_info=True,
@@ -389,7 +394,7 @@ class MitMDatasetRestApi(BaseSupersetModelRestApi):
def bulk_delete(self, **kwargs: Any) -> Response:
item_ids = kwargs['rison']
try:
- DeleteMitMDatasetCommand(item_ids).run()
+ DeleteMitMDatasetsCommand(item_ids).run()
return self.response(
200,
message=_(
@@ -405,6 +410,35 @@ class MitMDatasetRestApi(BaseSupersetModelRestApi):
except MitMDatasetDeleteFailedError as ex:
return self.response_422(message=str(ex))
+ @expose('/delete_with_related/', methods=('DELETE',))
+ @protect()
+ @safe
+ @statsd_metrics
+ @rison(get_delete_ids_schema)
+ @event_logger.log_this_with_context(
+ action=lambda self, *args,
+ **kwargs: f'{self.__class__.__name__}.bulk_delete_with_related',
+ log_to_statsd=False,
+ )
+ def bulk_delete_with_related(self, **kwargs: Any) -> Response:
+ item_ids = kwargs['rison']
+ try:
+ DeleteMitMDatasetsCommand(item_ids, delete_related=True).run()
+ return self.response(
+ 200,
+ message=_(
+ 'Deleted %(num)d dataset with its related objects',
+ 'Deleted %(num)d datasets along with their related objects',
+ num=len(item_ids),
+ ),
+ )
+ except MitMDatasetNotFoundError:
+ return self.response_404()
+ except MitMDatasetForbiddenError:
+ return self.response_403()
+ except MitMDatasetDeleteFailedError as ex:
+ return self.response_422(message=str(ex))
+
@expose('/export/', methods=('GET',))
@protect()
@safe
diff --git a/superset/tasks/mitm/__init__.py b/superset/tasks/mitm/__init__.py
index e69de29bb2..0e0c5a1eba 100644
--- a/superset/tasks/mitm/__init__.py
+++ b/superset/tasks/mitm/__init__.py
@@ -0,0 +1,4 @@
+from .call_external_service import call_external_service_task
+from .download_mitm_dataset import download_mitm_dataset_task
+from .drop_mitm_dataset import drop_mitm_dataset_task
+from .upload_mitm_dataset import upload_mitm_dataset_task
diff --git a/superset/tasks/mitm/call_external_service.py b/superset/tasks/mitm/call_external_service.py
index 16e91197eb..0483eac2eb 100644
--- a/superset/tasks/mitm/call_external_service.py
+++ b/superset/tasks/mitm/call_external_service.py
@@ -28,6 +28,7 @@ from superset.utils.cache import generate_cache_key
from superset.utils.core import override_user
from .common import service_call_timeout
from ..async_queries import _load_user_from_job_metadata
+from superset.extensions import celery_app
logger = logging.getLogger(__name__)
@@ -37,7 +38,7 @@ if TYPE_CHECKING:
CompleteForwardableRequest
-@shared_task(name='call_external_service',
+@celery_app.task(name='call_external_service',
soft_time_limit=service_call_timeout,
ignore_result=True,
pydantic=True)
diff --git a/superset/tasks/mitm/download_mitm_dataset.py b/superset/tasks/mitm/download_mitm_dataset.py
index b1cd2f923a..541f4440f7 100644
--- a/superset/tasks/mitm/download_mitm_dataset.py
+++ b/superset/tasks/mitm/download_mitm_dataset.py
@@ -3,9 +3,9 @@ from __future__ import annotations
import logging
from typing import TYPE_CHECKING
-from celery import shared_task
from celery.exceptions import SoftTimeLimitExceeded
+from superset.extensions import celery_app
from superset.utils.core import override_user
from .common import *
from ..async_queries import _load_user_from_job_metadata
@@ -19,9 +19,9 @@ if TYPE_CHECKING:
from superset.customization.external_service_support.common import AsyncJobMetadata
-@shared_task(name='download_mitm_dataset',
- soft_time_limit=service_call_timeout,
- ignore_result=True)
+@celery_app.task(name='download_mitm_dataset',
+ soft_time_limit=service_call_timeout,
+ ignore_result=True)
def download_mitm_dataset_task(job_metadata: AsyncJobMetadata,
mitm_dataset_id: int, result_base_url: str) -> None:
external_service_call_manager.call_started(job_metadata)
diff --git a/superset/tasks/mitm/drop_mitm_dataset.py b/superset/tasks/mitm/drop_mitm_dataset.py
index 09b536019c..2a2062d122 100644
--- a/superset/tasks/mitm/drop_mitm_dataset.py
+++ b/superset/tasks/mitm/drop_mitm_dataset.py
@@ -3,9 +3,9 @@ from __future__ import annotations
import logging
from typing import TYPE_CHECKING
-from celery import shared_task
from celery.exceptions import SoftTimeLimitExceeded
+from superset.extensions import celery_app
from superset.utils.core import override_user
from .common import *
from ..async_queries import _load_user_from_job_metadata
@@ -18,11 +18,11 @@ if TYPE_CHECKING:
from superset.customization.external_service_support.common import AsyncJobMetadata
-@shared_task(name='drop_mitm_dataset',
- soft_time_limit=service_call_timeout,
- ignore_result=True)
-def drop_mitm_dataset_task(job_metadata: AsyncJobMetadata, mitm_dataset_id: int) -> None:
-
+@celery_app.task(name='drop_mitm_dataset',
+ soft_time_limit=service_call_timeout,
+ ignore_result=True)
+def drop_mitm_dataset_task(job_metadata: AsyncJobMetadata,
+ mitm_dataset_id: int) -> None:
external_service_call_manager.call_started(job_metadata)
from ...commands.mitm.mitm_service.delete import DeleteUploadedMitMDatasetCommand
diff --git a/superset/tasks/mitm/upload_mitm_dataset.py b/superset/tasks/mitm/upload_mitm_dataset.py
index 074407d6da..1009f1e15a 100644
--- a/superset/tasks/mitm/upload_mitm_dataset.py
+++ b/superset/tasks/mitm/upload_mitm_dataset.py
@@ -3,9 +3,9 @@ from __future__ import annotations
import logging
from typing import TYPE_CHECKING
-from celery import shared_task
from celery.exceptions import SoftTimeLimitExceeded
+from superset.extensions import celery_app
from superset.utils.core import override_user
from .common import *
from ..async_queries import _load_user_from_job_metadata
@@ -18,9 +18,9 @@ if TYPE_CHECKING:
from superset.customization.external_service_support.common import AsyncJobMetadata
-@shared_task(name='upload_mitm_dataset',
- soft_time_limit=service_call_timeout,
- ignore_result=True)
+@celery_app.task(name='upload_mitm_dataset',
+ soft_time_limit=service_call_timeout,
+ ignore_result=True)
def upload_mitm_dataset_task(job_metadata: AsyncJobMetadata,
dataset_name: str, mitm_zip: bytes,
mitm_name: str = 'MAED') -> None:
diff --git a/tests/requests/mitm_dataset_api.http b/tests/requests/mitm_dataset_api.http
index e7150e771c..e2ffa92afc 100644
--- a/tests/requests/mitm_dataset_api.http
+++ b/tests/requests/mitm_dataset_api.http
@@ -74,3 +74,12 @@ POST http://localhost:8088/api/v1/mitm_dataset/export/
Authorization: Bearer {{jwt_token}}
X-CSRFToken: {{csrf_token}}
+###
+
+// @name Get Specific MitM Dataset
+GET http://localhost:8088/api/v1/mitm_dataset/d4a65922-f083-441f-be9d-3b7c88b283c9
+Authorization: Bearer {{jwt_token}}
+X-CSRFToken: {{csrf_token}}
+
+###
+
--
GitLab