diff --git a/src/preset_cli/api/clients/superset.py b/src/preset_cli/api/clients/superset.py index 3d6af95b..9853d4cd 100644 --- a/src/preset_cli/api/clients/superset.py +++ b/src/preset_cli/api/clients/superset.py @@ -296,7 +296,7 @@ def get_data( # pylint: disable=too-many-locals, too-many-arguments """ Run a dimensional query. """ - dataset = self.get_dataset(dataset_id)["result"] + dataset = self.get_dataset(dataset_id) if time_column is None: time_columns = [ @@ -394,7 +394,7 @@ def get_resource(self, resource_name: str, resource_id: int) -> Any: response = self.session.get(url) validate_response(response) - resource = response.json() + resource = response.json()["result"] return resource diff --git a/src/preset_cli/cli/superset/sync/dbt/command.py b/src/preset_cli/cli/superset/sync/dbt/command.py index 73a7b615..521468c2 100644 --- a/src/preset_cli/cli/superset/sync/dbt/command.py +++ b/src/preset_cli/cli/superset/sync/dbt/command.py @@ -303,7 +303,7 @@ def dbt_cloud( # pylint: disable=too-many-arguments, too-many-locals raise Exception("More than one database with the same name found") # need to get the database by itself so the response has the SQLAlchemy URI - database = superset_client.get_database(databases[0]["id"])["result"] + database = superset_client.get_database(databases[0]["id"]) models = dbt_client.get_models(job_id) models = apply_select(models, select, exclude) diff --git a/src/preset_cli/cli/superset/sync/dbt/exposures.py b/src/preset_cli/cli/superset/sync/dbt/exposures.py index c29dc248..8b145ec2 100644 --- a/src/preset_cli/cli/superset/sync/dbt/exposures.py +++ b/src/preset_cli/cli/superset/sync/dbt/exposures.py @@ -34,7 +34,7 @@ def get_chart_depends_on( query_context = json.loads(chart["query_context"]) dataset_id = query_context["datasource"]["id"] - dataset = client.get_dataset(dataset_id)["result"] + dataset = client.get_dataset(dataset_id) extra = json.loads(dataset["extra"] or "{}") if "depends_on" in extra: return [extra["depends_on"]] @@ -66,7 +66,7 @@ def get_dashboard_depends_on( depends_on = [] for dataset in payload["result"]: - full_dataset = client.get_dataset(int(dataset["id"]))["result"] + full_dataset = client.get_dataset(int(dataset["id"])) try: extra = json.loads(full_dataset["extra"] or "{}") except json.decoder.JSONDecodeError: @@ -114,7 +114,7 @@ def sync_exposures( # pylint: disable=too-many-locals dashboards_ids.add(dashboard["id"]) for chart_id in charts_ids: - chart = client.get_chart(chart_id)["result"] + chart = client.get_chart(chart_id) first_owner = chart["owners"][0] exposure = { "name": chart["slice_name"] + " [chart]", @@ -135,7 +135,7 @@ def sync_exposures( # pylint: disable=too-many-locals exposures.append(exposure) for dashboard_id in dashboards_ids: - dashboard = client.get_dashboard(dashboard_id)["result"] + dashboard = client.get_dashboard(dashboard_id) first_owner = dashboard["owners"][0] exposure = { "name": dashboard["dashboard_title"] + " [dashboard]", diff --git a/tests/api/clients/superset_test.py b/tests/api/clients/superset_test.py index a6c45385..b2fa291f 100644 --- a/tests/api/clients/superset_test.py +++ b/tests/api/clients/superset_test.py @@ -695,7 +695,7 @@ def test_get_resource(requests_mock: Mocker) -> None: # the payload schema is irrelevant, since it's passed through unmodified requests_mock.get( "https://superset.example.org/api/v1/database/1", - json={"Hello": "world"}, + json={"result": {"Hello": "world"}}, ) auth = Auth() diff --git a/tests/cli/superset/sync/dbt/command_test.py b/tests/cli/superset/sync/dbt/command_test.py index 9ffbb341..630796df 100644 --- a/tests/cli/superset/sync/dbt/command_test.py +++ b/tests/cli/superset/sync/dbt/command_test.py @@ -462,7 +462,7 @@ def test_dbt_cloud(mocker: MockerFixture) -> None: dbt_client.get_metrics.return_value = metrics database = mocker.MagicMock() superset_client.get_databases.return_value = [database] - superset_client.get_database.return_value = {"result": database} + superset_client.get_database.return_value = database runner = CliRunner() result = runner.invoke( @@ -532,7 +532,7 @@ def test_dbt_cloud_no_job_id(mocker: MockerFixture) -> None: dbt_client.get_jobs.return_value = [{"id": 123, "name": "My job"}] database = mocker.MagicMock() superset_client.get_databases.return_value = [database] - superset_client.get_database.return_value = {"result": database} + superset_client.get_database.return_value = database runner = CliRunner() result = runner.invoke( diff --git a/tests/cli/superset/sync/dbt/exposures_test.py b/tests/cli/superset/sync/dbt/exposures_test.py index 7a16175d..6e9f69d4 100644 --- a/tests/cli/superset/sync/dbt/exposures_test.py +++ b/tests/cli/superset/sync/dbt/exposures_test.py @@ -500,7 +500,7 @@ def test_get_dashboard_depends_on(mocker: MockerFixture) -> None: Test ``get_dashboard_depends_on``. """ client = mocker.MagicMock() - client.get_dataset.return_value = dataset_response + client.get_dataset.return_value = dataset_response["result"] session = client.auth.session session.get().json.return_value = datasets_response @@ -515,7 +515,7 @@ def test_get_dashboard_depends_on_no_extra(mocker: MockerFixture) -> None: client = mocker.MagicMock() modified_dataset_response = copy.deepcopy(dataset_response) modified_dataset_response["result"]["extra"] = None # type: ignore - client.get_dataset.return_value = modified_dataset_response + client.get_dataset.return_value = modified_dataset_response["result"] session = client.auth.session session.get().json.return_value = datasets_response @@ -530,7 +530,7 @@ def test_get_dashboard_depends_on_invalid_extra(mocker: MockerFixture) -> None: client = mocker.MagicMock() modified_dataset_response = copy.deepcopy(dataset_response) modified_dataset_response["result"]["extra"] = "{[(" # type: ignore - client.get_dataset.return_value = modified_dataset_response + client.get_dataset.return_value = modified_dataset_response["result"] session = client.auth.session session.get().json.return_value = datasets_response @@ -543,7 +543,7 @@ def test_get_chart_depends_on(mocker: MockerFixture) -> None: Test ``get_chart_depends_on``. """ client = mocker.MagicMock() - client.get_dataset.return_value = dataset_response + client.get_dataset.return_value = dataset_response["result"] depends_on = get_chart_depends_on(client, chart_response["result"], {}) assert depends_on == ["ref('messages_channels')"] @@ -556,7 +556,7 @@ def test_get_chart_depends_on_no_extra(mocker: MockerFixture) -> None: client = mocker.MagicMock() modified_dataset_response = copy.deepcopy(dataset_response) modified_dataset_response["result"]["extra"] = None # type: ignore - client.get_dataset.return_value = modified_dataset_response + client.get_dataset.return_value = modified_dataset_response["result"] depends_on = get_chart_depends_on(client, chart_response["result"], {}) assert not depends_on @@ -572,8 +572,8 @@ def test_sync_exposures(mocker: MockerFixture, fs: FakeFilesystem) -> None: client = mocker.MagicMock() client.baseurl = URL("https://superset.example.org/") - client.get_chart.return_value = chart_response - client.get_dashboard.return_value = dashboard_response + client.get_chart.return_value = chart_response["result"] + client.get_dashboard.return_value = dashboard_response["result"] session = client.auth.session session.get().json.return_value = related_objects_response mocker.patch( @@ -658,7 +658,7 @@ def test_get_chart_depends_on_from_dataset(mocker: MockerFixture) -> None: client = mocker.MagicMock() modified_dataset_response = copy.deepcopy(dataset_response) modified_dataset_response["result"]["extra"] = None # type: ignore - client.get_dataset.return_value = modified_dataset_response + client.get_dataset.return_value = modified_dataset_response["result"] key = ModelKey("public", "messages_channels") depends_on = get_chart_depends_on( @@ -676,7 +676,7 @@ def test_get_dashboard_depends_on_from_dataset(mocker: MockerFixture) -> None: client = mocker.MagicMock() modified_dataset_response = copy.deepcopy(dataset_response) modified_dataset_response["result"]["extra"] = None # type: ignore - client.get_dataset.return_value = modified_dataset_response + client.get_dataset.return_value = modified_dataset_response["result"] session = client.auth.session session.get().json.return_value = datasets_response