[metastore] add config flag to fetch sample data for views (#4172)

- This commit introduces a new configuration flag, allow_sample_data_from_views, to control whether Hue attempts to fetch sample data for database views.

- Previously, fetching sample data from views was disabled by default due to potential performance concerns with complex or long-running views. This change makes this behavior configurable.

- Default value: false (maintains previous behavior)

- If set to true, Hue will attempt to fetch sample data for views.

---------

Co-authored-by: Harsh Gupta <[email protected]>

Author: bjornalm

apps/beeswax/src/beeswax/api.py

@@ -15,9 +15,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import re
 import json
 import logging
+import re
 from builtins import zip
 
 from django.http import Http404
@@ -33,7 +33,7 @@
 from beeswax.forms import QueryForm
 from beeswax.models import QueryHistory, Session
 from beeswax.server import dbms
-from beeswax.server.dbms import QueryServerException, QueryServerTimeoutException, SubQueryTable, expand_exception, get_query_server_config
+from beeswax.server.dbms import expand_exception, get_query_server_config, QueryServerException, QueryServerTimeoutException, SubQueryTable
 from beeswax.views import (
   _get_query_handle_and_state,
   authorized_get_design,
@@ -52,9 +52,9 @@
 from desktop.lib.i18n import force_unicode
 from desktop.lib.parameterization import substitute_variables
 from metastore import parser
-from metastore.conf import FORCE_HS2_METADATA
+from metastore.conf import ALLOW_SAMPLE_DATA_FROM_VIEWS, FORCE_HS2_METADATA
 from metastore.views import _get_db, _get_servername
-from notebook.models import MockedDjangoRequest, escape_rows, make_notebook
+from notebook.models import escape_rows, make_notebook, MockedDjangoRequest
 from useradmin.models import User
 
 LOG = logging.getLogger()
@@ -606,7 +606,7 @@ def save_results_hdfs_file(request, query_history_id):
 
       try:
         handle, state = _get_query_handle_and_state(query_history)
-      except Exception as ex:
+      except Exception:
         response['message'] = _('Cannot find query handle and state: %s') % str(query_history)
         response['status'] = -2
         return JsonResponse(response)
@@ -669,7 +669,7 @@ def save_results_hive_table(request, query_history_id):
       try:
         handle, state = _get_query_handle_and_state(query_history)
         result_meta = db.get_results_metadata(handle)
-      except Exception as ex:
+      except Exception:
         response['message'] = _('Cannot find query handle and state: %s') % str(query_history)
         response['status'] = -2
         return JsonResponse(response)
@@ -733,7 +733,7 @@ def _get_sample_data(db, database, table, column, nested, is_async=False, cluste
       query_server = get_query_server_config('impala', connector=cluster)
       db = dbms.get(db.client.user, query_server, cluster=cluster)
 
-  if table_obj and table_obj.is_view:
+  if table_obj and table_obj.is_view and not ALLOW_SAMPLE_DATA_FROM_VIEWS.get():
     response = {'status': -1}
     response['message'] = _('Not getting sample data as this is a view which can be expensive when run.')
     return response

apps/beeswax/src/beeswax/api_tests.py

@@ -16,25 +16,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import sys
 import logging
 from unittest.mock import Mock, patch
 
 import pytest
-from django.test import TestCase
 from requests.exceptions import ReadTimeout
 
-from beeswax.api import _autocomplete, get_functions
+from beeswax.api import _autocomplete, _get_functions, _get_sample_data
 from desktop.lib.django_test_util import make_logged_in_client
-from desktop.lib.test_utils import add_to_group, grant_access
+from metastore.conf import ALLOW_SAMPLE_DATA_FROM_VIEWS
 from useradmin.models import User
 
 LOG = logging.getLogger()
 
 
 @pytest.mark.django_db
-class TestApi():
-
+class TestApi:
   def setup_method(self):
     self.client = make_logged_in_client(username="test", groupname="default", recreate=True, is_superuser=False)
     self.user = User.objects.get(username="test")
@@ -43,75 +40,99 @@ def test_autocomplete_time_out(self):
     get_tables_meta = Mock(
       side_effect=ReadTimeout("HTTPSConnectionPool(host='gethue.com', port=10001): Read timed out. (read timeout=120)")
     )
-    db = Mock(
-      get_tables_meta=get_tables_meta
-    )
+    db = Mock(get_tables_meta=get_tables_meta)
 
-    resp = _autocomplete(db, database='database')
+    resp = _autocomplete(db, database="database")
 
-    assert (
-      resp ==
-      {
-        'code': 500,
-        'error': "HTTPSConnectionPool(host='gethue.com', port=10001): Read timed out. (read timeout=120)"
-      })
+    assert resp == {"code": 500, "error": "HTTPSConnectionPool(host='gethue.com', port=10001): Read timed out. (read timeout=120)"}
 
   def test_get_functions(self):
-    db = Mock(
-      get_functions=Mock(
-        return_value=Mock(
-          rows=Mock(
-            return_value=[{'name': 'f1'}, {'name': 'f2'}]
-          )
-        )
-      )
-    )
+    # Mock db.get_functions() to return rows that escape_rows can process
+    # Each row should be a list where row[0] is the function name
+    db = Mock()
+    db.get_functions = Mock(return_value=[["f1"], ["f2"]])  # Return list of rows
+    db.client = Mock(query_server={"dialect": "hive"})  # Non-Impala dialect
 
-    resp = get_functions(db)
+    resp = _get_functions(db)  # Call the internal function
 
-    assert (
-      resp ==
-      [{'name': 'f1'}, {'name': 'f2'}])
+    assert resp == [{"name": "f1"}, {"name": "f2"}]
 
-  def test_get_functions(self):
-    with patch('beeswax.api._get_functions') as _get_functions:
+  def test_autocomplete_functions(self):
+    with patch("beeswax.api._get_functions") as _get_functions:
       db = Mock()
-      _get_functions.return_value = [
-        {'name': 'f1'}, {'name': 'f2'}, {'name': 'f3'}
-      ]
+      _get_functions.return_value = [{"name": "f1"}, {"name": "f2"}, {"name": "f3"}]
 
-      resp = _autocomplete(db, database='default', operation='functions')
+      resp = _autocomplete(db, database="default", operation="functions")
 
-      assert (
-        resp['functions'] ==
-        [{'name': 'f1'}, {'name': 'f2'}, {'name': 'f3'}])
+      assert resp["functions"] == [{"name": "f1"}, {"name": "f2"}, {"name": "f3"}]
 
   def test_get_function(self):
     db = Mock()
-    db.client = Mock(query_server={'dialect': 'hive'})
+    db.client = Mock(query_server={"dialect": "hive"})
     db.get_function = Mock(
       return_value=[
-        ['floor_month(param) - Returns the timestamp at a month granularity'],
-        ['param needs to be a timestamp value'],
-        ['Example:'],
+        ["floor_month(param) - Returns the timestamp at a month granularity"],
+        ["param needs to be a timestamp value"],
+        ["Example:"],
         ["> SELECT floor_month(CAST('yyyy-MM-dd HH:mm:ss' AS TIMESTAMP)) FROM src;"],
-        ['yyyy-MM-01 00:00:00']
+        ["yyyy-MM-01 00:00:00"],
       ]
     )
 
-    data = _autocomplete(db, database='floor_month', operation='function')
-
-    assert (
-      data['function'] ==
-      {
-        'name': 'floor_month',
-        'signature': 'floor_month(param)',
-        'description':
-            'Returns the timestamp at a month granularity\nparam needs to be a timestamp value\nExample:\n'
-            '> SELECT floor_month(CAST(\'yyyy-MM-dd HH:mm:ss\' AS TIMESTAMP)) FROM src;\nyyyy-MM-01 00:00:00'
-      })
-
-    db.client = Mock(query_server={'dialect': 'impala'})
-    data = _autocomplete(db, operation='function')
-
-    assert data['function'] == {}
+    data = _autocomplete(db, database="floor_month", operation="function")
+
+    assert data["function"] == {
+      "name": "floor_month",
+      "signature": "floor_month(param)",
+      "description": "Returns the timestamp at a month granularity\nparam needs to be a timestamp value\nExample:\n"
+      "> SELECT floor_month(CAST('yyyy-MM-dd HH:mm:ss' AS TIMESTAMP)) FROM src;\nyyyy-MM-01 00:00:00",
+    }
+
+    db.client = Mock(query_server={"dialect": "impala"})
+    data = _autocomplete(db, operation="function")
+
+    assert data["function"] == {}
+
+  @patch("beeswax.api.dbms.get")
+  def test_get_sample_data_for_views(self, mock_dbms_get):
+    # Mock table_obj
+    table_obj_mock = Mock(is_view=True, is_impala_only=False)
+
+    # Mock the db object that dbms.get() would return
+    db_mock = Mock(get_table=Mock(return_value=table_obj_mock))
+    mock_dbms_get.return_value = db_mock
+
+    # Scenario 1: allow_sample_data_from_views is False
+    reset = ALLOW_SAMPLE_DATA_FROM_VIEWS.set_for_testing(False)
+    try:
+      response = _get_sample_data(db_mock, "default_db", "test_view_table", None, None)
+
+      assert response == {
+        "status": -1,
+        "message": "Not getting sample data as this is a view which can be expensive when run.",
+      }
+    finally:
+      reset()
+
+    # Scenario 2: allow_sample_data_from_views is True
+    reset = ALLOW_SAMPLE_DATA_FROM_VIEWS.set_for_testing(True)
+    try:
+      # Mock db.get_sample to simulate successful data fetching past the view check
+      # We expect it to be called if the view check is passed.
+      db_mock.get_sample.return_value = Mock(
+        rows=Mock(return_value=[["col1_val", "col2_val"]]),
+        cols=Mock(return_value=["col1", "col2"]),
+        full_cols=Mock(return_value=[{"name": "col1"}, {"name": "col2"}]),
+      )
+      mock_dbms_get.return_value = db_mock
+
+      response = _get_sample_data(db_mock, "default_db", "test_view_table", None, None)
+      assert response == {
+        "status": 0,
+        "headers": ["col1", "col2"],
+        "full_headers": [{"name": "col1"}, {"name": "col2"}],
+        "rows": [["col1_val", "col2_val"]],
+      }
+      db_mock.get_sample.assert_called_once_with("default_db", table_obj_mock, None, None, generate_sql_only=False, operation=None)
+    finally:
+      reset()

apps/beeswax/src/beeswax/conf.py

@@ -15,19 +15,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import math
 import logging
+import math
 import os.path
 
-from django.utils.translation import gettext as _, gettext_lazy as _t
+from django.utils.translation import gettext_lazy as _t
 
 from desktop.conf import (
   AUTH_PASSWORD as DEFAULT_AUTH_PASSWORD,
   AUTH_USERNAME as DEFAULT_AUTH_USERNAME,
   default_ssl_cacerts,
   default_ssl_validate,
 )
-from desktop.lib.conf import Config, ConfigSection, coerce_bool, coerce_csv, coerce_password_from_script
+from desktop.lib.conf import coerce_bool, coerce_csv, coerce_password_from_script, Config, ConfigSection
 
 LOG = logging.getLogger()
 

apps/metastore/src/metastore/conf.py

@@ -17,7 +17,7 @@
 
 from django.utils.translation import gettext_lazy as _
 
-from desktop.lib.conf import Config, coerce_bool
+from desktop.lib.conf import coerce_bool, Config
 
 ENABLE_NEW_CREATE_TABLE = Config(
   key="enable_new_create_table",
@@ -41,3 +41,10 @@
   type=coerce_bool,
   help=_('Choose whether to show the table ERD component.')
 )
+
+ALLOW_SAMPLE_DATA_FROM_VIEWS = Config(
+  key='allow_sample_data_from_views',
+  default=False,
+  type=coerce_bool,
+  help=_('Choose whether to allow fetching sample data from views. By default, this is false to prevent potentially expensive queries.')
+)

desktop/conf.dist/hue.ini

@@ -1578,6 +1578,10 @@ submit_to=True
 # Choose whether to show the table ERD component. Default false
 ## show_table_erd=false
 
+# Choose whether to allow fetching sample data from views.
+# By default, this is false to prevent potentially expensive queries.
+## allow_sample_data_from_views=false
+
 ###########################################################################
 # Settings to configure Impala
 ###########################################################################

desktop/conf/pseudo-distributed.ini.tmpl

@@ -1547,7 +1547,6 @@
     # Choose whether Hue should validate certificates received from the server.
     ## validate=true
 
-
 ###########################################################################
 # Settings to configure Metastore
 ###########################################################################
@@ -1556,6 +1555,9 @@
   # Flag to turn on the new version of the create table wizard.
   ## enable_new_create_table=true
 
+  # Allow fetching sample data from views. By default, this is false to prevent potentially expensive queries.
+  ## allow_sample_data_from_views=false
+
   # Flag to force all metadata calls (e.g. list tables, table or column details...) to happen via HiveServer2 if available instead of Impala.
   ## force_hs2_metadata=false
 

desktop/core/src/desktop/api2.py

@@ -92,6 +92,7 @@
   search_entities_interactive as metadata_search_entities_interactive,
 )
 from metadata.conf import has_catalog
+from metastore.conf import ALLOW_SAMPLE_DATA_FROM_VIEWS
 from notebook.connectors.base import get_interpreter, Notebook
 from notebook.management.commands import notebook_setup
 from pig.management.commands import pig_setup
@@ -146,6 +147,7 @@ def get_config(request):
     'is_yarn_enabled': is_yarn(),
     'enable_task_server': TASK_SERVER_V2.ENABLED.get(),
     'enable_workflow_creation_action': ENABLE_WORKFLOW_CREATION_ACTION.get(),
+    'allow_sample_data_from_views': ALLOW_SAMPLE_DATA_FROM_VIEWS.get(),
   }
 
   # Storage browser configuration

desktop/core/src/desktop/js/catalog/DataCatalogEntry.ts

@@ -18,6 +18,7 @@ import * as ko from 'knockout';
 import KnockoutObservable from '@types/knockout';
 
 import { Cancellable, CancellablePromise } from 'api/cancellablePromise';
+import { getLastKnownConfig } from 'config/hueConfig';
 import {
   addNavTags,
   deleteNavTags,
@@ -1662,7 +1663,8 @@ export default class DataCatalogEntry {
       operation?: string;
     }
   ): CancellablePromise<Sample> {
-    if (this.isView()) {
+    const config = getLastKnownConfig();
+    if (this.isView() && (!config || !config.hue_config?.allow_sample_data_from_views)) {
       return CancellablePromise.reject();
     }
 

desktop/core/src/desktop/js/config/types.ts

@@ -96,6 +96,7 @@ export interface HueConfig extends GenericApiResponse {
     enable_task_server: boolean;
     is_admin: boolean;
     is_yarn_enabled: boolean;
+    allow_sample_data_from_views: boolean;
   };
   storage_browser: StorageBrowserConfig;
   hue_version?: string;