diff --git a/desktop/libs/dashboard/src/dashboard/views.py b/desktop/libs/dashboard/src/dashboard/views.py index 3128c53f5fe..2c01ba38db5 100644 --- a/desktop/libs/dashboard/src/dashboard/views.py +++ b/desktop/libs/dashboard/src/dashboard/views.py @@ -147,7 +147,9 @@ def new_search(request): engine = request.GET.get('engine', 'solr') cluster = request.POST.get('cluster','""') - collections = get_engine(request.user, engine, cluster=cluster).datasets() if engine != 'report' else ['default'] + spark_refresh_token = {'SPARK_REFRESH_TOKEN': request.session.get('oidc_refresh_token')} + collections = get_engine(request.user, engine, cluster=cluster).\ + datasets(refresh_token=spark_refresh_token) if engine != 'report' else ['default'] if not collections: if engine == 'solr': diff --git a/desktop/libs/notebook/src/notebook/dashboard_api.py b/desktop/libs/notebook/src/notebook/dashboard_api.py index fc5bb1ebc5e..6c18502290c 100644 --- a/desktop/libs/notebook/src/notebook/dashboard_api.py +++ b/desktop/libs/notebook/src/notebook/dashboard_api.py @@ -34,6 +34,7 @@ from notebook.models import make_notebook from notebook.connectors.base import get_api, OperationTimeout, Notebook +from notebook.connectors.hiveserver2 import HS2Api from dashboard.dashboard_api import DashboardApi from dashboard.models import Collection2, augment_response @@ -237,12 +238,21 @@ def fetch_result(self, dashboard, query, facet): # This method currently behaves more like a static method - def datasets(self, show_all=False, database=None): + def datasets(self, show_all=False, database=None, refresh_token: dict = None): snippet = {'type': self.engine} # Ideally from left assist at some point instead if database is None: - databases = get_api(MockRequest(self.user, self.cluster), snippet).autocomplete(snippet)['databases'] + api = get_api(MockRequest(self.user, self.cluster), snippet) + if isinstance(api, HS2Api): + databases = api.autocomplete(snippet, refresh_token=refresh_token)['databases'] + database = databases and 'default' not in databases and sorted(databases)[0] or 'default' + return [ + database + '.' + table['name'] + for table in api.autocomplete(snippet, database=database, refresh_token=refresh_token)['tables_meta'] + ] + else: + databases = api.autocomplete(snippet)['databases'] database = databases and 'default' not in databases and sorted(databases)[0] or 'default' return [