Copilot commented on code in PR #36550: URL: https://github.com/apache/superset/pull/36550#discussion_r2665962370
########## tests/unit_tests/views/datasource/utils_test.py: ########## @@ -0,0 +1,217 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Tests for superset.views.datasource.utils module.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from superset.errors import ErrorLevel, SupersetError, SupersetErrorType +from superset.exceptions import SupersetSecurityException + + +@patch("superset.views.datasource.utils.get_limit_clause") +def test_get_samples_raises_security_exception_when_access_denied( + mock_get_limit_clause: MagicMock, +): + """ + Test that get_samples() enforces access control by calling raise_for_access(). + This verifies the fix for issue #31944 where users with "can samples on Datasource" + permission could read samples from datasets they don't have access to. + """ + mock_get_limit_clause.return_value = {"row_offset": 0, "row_limit": 100} + + mock_datasource = MagicMock() + mock_datasource.type = "table" + mock_datasource.id = 1 + mock_datasource.columns = [] + + mock_samples_context = MagicMock() + mock_count_context = MagicMock() + + # Simulate security exception when raise_for_access is called + mock_samples_context.raise_for_access.side_effect = SupersetSecurityException( + SupersetError( + message="Access denied", + error_type=SupersetErrorType.DATASOURCE_SECURITY_ACCESS_ERROR, + level=ErrorLevel.WARNING, + ) + ) + + with ( + patch( + "superset.views.datasource.utils.DatasourceDAO.get_datasource", + return_value=mock_datasource, + ), + patch( + "superset.views.datasource.utils.QueryContextFactory" + ) as mock_factory_class, + ): + mock_factory = MagicMock() + mock_factory_class.return_value = mock_factory + + # Return different mock contexts for samples vs count queries + mock_factory.create.side_effect = [mock_samples_context, mock_count_context] + + from superset.views.datasource.utils import get_samples + + with pytest.raises(SupersetSecurityException) as exc_info: + get_samples( + datasource_type="table", + datasource_id=1, + force=False, + page=1, + per_page=100, + ) + + assert exc_info.value.error.error_type == ( + SupersetErrorType.DATASOURCE_SECURITY_ACCESS_ERROR + ) + # Verify raise_for_access was called on the samples context + mock_samples_context.raise_for_access.assert_called_once() + + +@patch("superset.views.datasource.utils.get_limit_clause") +def test_get_samples_calls_raise_for_access_on_both_contexts( + mock_get_limit_clause: MagicMock, +): + """ + Test that get_samples() calls raise_for_access() on both the samples + and count_star query contexts before fetching data. + """ + mock_get_limit_clause.return_value = {"row_offset": 0, "row_limit": 100} + + mock_datasource = MagicMock() + mock_datasource.type = "table" + mock_datasource.id = 1 + mock_datasource.columns = [] + + mock_samples_context = MagicMock() + mock_count_context = MagicMock() + + # Set up successful access check + mock_samples_context.raise_for_access.return_value = None + mock_count_context.raise_for_access.return_value = None + + # Set up successful payload responses + mock_count_context.get_payload.return_value = { + "queries": [{"data": [{"COUNT(*)": 100}], "status": "success"}] + } + mock_samples_context.get_payload.return_value = { + "queries": [ + { + "data": [{"col1": "val1"}], + "status": "success", + "cache_key": "test_key", + } + ] + } + + with ( + patch( + "superset.views.datasource.utils.DatasourceDAO.get_datasource", + return_value=mock_datasource, + ), + patch( + "superset.views.datasource.utils.QueryContextFactory" + ) as mock_factory_class, + ): + mock_factory = MagicMock() + mock_factory_class.return_value = mock_factory + + # Return different mock contexts for samples vs count queries + mock_factory.create.side_effect = [mock_samples_context, mock_count_context] + + from superset.views.datasource.utils import get_samples + + result = get_samples( + datasource_type="table", + datasource_id=1, + force=False, + page=1, + per_page=100, + ) + + # Verify both contexts had raise_for_access called + mock_samples_context.raise_for_access.assert_called_once() + mock_count_context.raise_for_access.assert_called_once() + + # Verify the result contains expected data + assert result["data"] == [{"col1": "val1"}] + assert result["total_count"] == 100 + Review Comment: The test suite should include a test case that verifies access control enforcement when the payload parameter is provided. When payload is not None, get_samples() follows a different code path (DRILL_DETAIL result type instead of SAMPLES), and this security-critical path should also be tested to ensure raise_for_access() is called correctly in both scenarios. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
