fix: dataset update with invalid SQL query (#35543)

tests/unit_tests/commands/dataset/update_test.py

@@ -70,6 +70,105 @@ def test_update_dataset_forbidden(mocker: MockerFixture) -> None:
         UpdateDatasetCommand(1, {"name": "test"}).run()
 
 
+def test_update_dataset_sql_authorized_schema(mocker: MockerFixture) -> None:
+    """
+    Test that updating a dataset with SQL works when user has schema access.
+    """
+    mock_dataset_dao = mocker.patch("superset.commands.dataset.update.DatasetDAO")
+    mock_database = mocker.MagicMock()
+    mock_database.id = 1
+    mock_database.get_default_catalog.return_value = "catalog"
+    mock_database.allow_multi_catalog = False
+
+    mock_dataset = mocker.MagicMock()
+    mock_dataset.database = mock_database
+    mock_dataset.catalog = "catalog"
+    mock_dataset.schema = "public"
+    mock_dataset.table_name = "test_table"
+    mock_dataset.owners = []  # No owners to avoid ownership computation issues
+
+    mock_dataset_dao.find_by_id.return_value = mock_dataset
+    mock_dataset_dao.get_database_by_id.return_value = mock_database
+    mock_dataset_dao.validate_update_uniqueness.return_value = True
+    mock_dataset_dao.update.return_value = mock_dataset
+
+    # Mock successful ownership check
+    mocker.patch(
+        "superset.commands.dataset.update.security_manager.raise_for_ownership",
+    )
+
+    # Mock security manager methods for owner computation
+    mocker.patch("superset.commands.utils.security_manager.is_admin", return_value=True)
+
+    # Mock security manager to allow access to the schema
+    mocker.patch(
+        "superset.commands.dataset.update.security_manager.raise_for_access",
+    )
+
+    # Update dataset with SQL - should work when user has access
+    result = UpdateDatasetCommand(
+        1, {"sql": "SELECT * FROM public.allowed_table"}
+    ).run()
+
+    # Verify the update was called
+    assert result == mock_dataset
+    mock_dataset_dao.update.assert_called_once()
+
+
+def test_update_dataset_sql_unauthorized_schema(mocker: MockerFixture) -> None:
+    """
+    Test that updating a dataset with SQL to an unauthorized schema raises an error.
+    """
+    mock_dataset_dao = mocker.patch("superset.commands.dataset.update.DatasetDAO")
+    mock_database = mocker.MagicMock()
+    mock_database.id = 1
+    mock_database.get_default_catalog.return_value = "catalog"
+    mock_database.allow_multi_catalog = False
+
+    mock_dataset = mocker.MagicMock()
+    mock_dataset.database = mock_database
+    mock_dataset.catalog = "catalog"
+    mock_dataset.schema = "public"
+    mock_dataset.table_name = "test_table"
+    mock_dataset.owners = []  # No owners to avoid ownership computation issues
+
+    mock_dataset_dao.find_by_id.return_value = mock_dataset
+    mock_dataset_dao.get_database_by_id.return_value = mock_database
+    mock_dataset_dao.validate_update_uniqueness.return_value = True
+
+    # Mock successful ownership check
+    mocker.patch(
+        "superset.commands.dataset.update.security_manager.raise_for_ownership",
+    )
+
+    # Mock security manager methods for owner computation
+    mocker.patch("superset.commands.utils.security_manager.is_admin", return_value=True)
+
+    # Mock security manager to raise error for SQL schema access
+    mocker.patch(
+        "superset.commands.dataset.update.security_manager.raise_for_access",
+        side_effect=SupersetSecurityException(
+            SupersetError(
+                error_type=SupersetErrorType.MISSING_OWNERSHIP_ERROR,
+                message="You don't have access to the 'restricted_schema' schema",
+                level=ErrorLevel.ERROR,
+            )
+        ),
+    )
+
+    # Try to update dataset with SQL querying an unauthorized schema
+    with pytest.raises(DatasetInvalidError) as excinfo:
+        UpdateDatasetCommand(
+            1, {"sql": "SELECT * FROM restricted_schema.sensitive_table"}
+        ).run()
+
+    # Check that the appropriate error message is in the exceptions
+    assert any(
+        "You don't have access to the 'restricted_schema' schema" in str(exc)
+        for exc in excinfo.value._exceptions
+    )
+
+
 @pytest.mark.parametrize(
     ("payload, exception, error_msg"),
     [