feat(sql lab): display presto and trino tracking url (#20799)

tests/integration_tests/conftest.py

@@ -16,15 +16,18 @@
 # under the License.
 from __future__ import annotations
 
+import contextlib
 import functools
+from operator import ge
 from typing import Any, Callable, Optional, TYPE_CHECKING
 from unittest.mock import patch
 
 import pytest
 from flask.ctx import AppContext
+from flask_appbuilder.security.sqla import models as ab_models
 from sqlalchemy.engine import Engine
 
-from superset import db
+from superset import db, security_manager
 from superset.extensions import feature_flag_manager
 from superset.utils.core import json_dumps_w_dates
 from superset.utils.database import get_example_database, remove_database
@@ -68,6 +71,50 @@ def login_as_admin(login_as: Callable[..., None]):
     yield login_as("admin")
 
 
+@pytest.fixture
+def create_user(app_context: AppContext):
+    def _create_user(username: str, role: str = "Admin", password: str = "general"):
+        security_manager.add_user(
+            username,
+            "firstname",
+            "lastname",
+            "email@exaple.com",
+            security_manager.find_role(role),
+            password,
+        )
+        return security_manager.find_user(username)
+
+    return _create_user
+
+
+@pytest.fixture
+def get_user(app_context: AppContext):
+    def _get_user(username: str) -> ab_models.User:
+        return (
+            db.session.query(security_manager.user_model)
+            .filter_by(username=username)
+            .one_or_none()
+        )
+
+    return _get_user
+
+
+@pytest.fixture
+def get_or_create_user(get_user, create_user) -> ab_models.User:
+    @contextlib.contextmanager
+    def _get_user(username: str) -> ab_models.User:
+        user = get_user(username)
+        if not user:
+            # if user is created by test, remove it after done
+            user = create_user(username)
+            yield user
+            db.session.delete(user)
+        else:
+            yield user
+
+    return _get_user
+
+
 @pytest.fixture(autouse=True, scope="session")
 def setup_sample_data() -> Any:
     # TODO(john-bodley): Determine a cleaner way of setting up the sample data without

tests/integration_tests/sql_lab/__init__.py

@@ -0,0 +1,16 @@
+# 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/integration_tests/sql_lab/conftest.py

@@ -0,0 +1,71 @@
+# 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.
+from typing import Callable, ContextManager
+
+import pytest
+from flask_appbuilder.security.sqla import models as ab_models
+
+from superset import db
+from superset.models.sql_lab import Query
+from superset.utils.core import shortid
+from superset.utils.database import get_example_database
+
+
+def force_async_run(allow_run_async: bool):
+    example_db = get_example_database()
+    orig_allow_run_async = example_db.allow_run_async
+
+    example_db.allow_run_async = allow_run_async
+    db.session.commit()
+
+    yield example_db
+
+    example_db.allow_run_async = orig_allow_run_async
+    db.session.commit()
+
+
+@pytest.fixture
+def non_async_example_db(app_context):
+    gen = force_async_run(False)
+    yield next(gen)
+    try:
+        next(gen)
+    except StopIteration:
+        pass
+
+
+@pytest.fixture
+def async_example_db(app_context):
+    gen = force_async_run(True)
+    yield next(gen)
+    try:
+        next(gen)
+    except StopIteration:
+        pass
+
+
+@pytest.fixture
+def example_query(get_or_create_user: Callable[..., ContextManager[ab_models.User]]):
+    with get_or_create_user("sqllab-test-user") as user:
+        query = Query(
+            client_id=shortid()[:10], database=get_example_database(), user=user
+        )
+        db.session.add(query)
+        db.session.commit()
+        yield query
+        db.session.delete(query)
+        db.session.commit()

tests/integration_tests/sql_lab/test_execute_sql_statements.py

@@ -0,0 +1,56 @@
+# 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.
+from superset import app, db
+from superset.common.db_query_status import QueryStatus
+from superset.models.core import Database
+from superset.models.sql_lab import Query
+from superset.sql_lab import execute_sql_statements
+from superset.utils.dates import now_as_float
+
+
+def test_non_async_execute(non_async_example_db: Database, example_query: Query):
+    """Test query.tracking_url is attached for Presto and Hive queries"""
+    result = execute_sql_statements(
+        example_query.id,
+        "select 1 as foo;",
+        store_results=False,
+        return_results=True,
+        session=db.session,
+        start_time=now_as_float(),
+        expand_data=True,
+        log_params=dict(),
+    )
+    assert result
+    assert result["query_id"] == example_query.id
+    assert result["status"] == QueryStatus.SUCCESS
+    assert result["data"] == [{"foo": 1}]
+
+    # should attach apply tracking URL for Presto & Hive
+    if non_async_example_db.db_engine_spec.engine == "presto":
+        assert example_query.tracking_url
+        assert "/ui/query.html?" in example_query.tracking_url
+
+        app.config["TRACKING_URL_TRANSFORMER"] = lambda url, query: url.replace(
+            "/ui/query.html?", f"/{query.client_id}/"
+        )
+        assert f"/{example_query.client_id}/" in example_query.tracking_url
+
+        app.config["TRACKING_URL_TRANSFORMER"] = lambda url: url + "&foo=bar"
+        assert example_query.tracking_url.endswith("&foo=bar")
+
+    if non_async_example_db.db_engine_spec.engine_name == "hive":
+        assert example_query.tracking_url_raw

tests/integration_tests/sqllab_tests.py

@@ -18,6 +18,7 @@
 """Unit tests for Sql Lab"""
 import json
 from datetime import datetime, timedelta
+from math import ceil, floor
 
 import pytest
 from celery.exceptions import SoftTimeLimitExceeded
@@ -70,8 +71,8 @@ class TestSqlLab(SupersetTestCase):
         db.session.query(Query).delete()
         db.session.commit()
         self.run_sql(QUERY_1, client_id="client_id_1", username="admin")
-        self.run_sql(QUERY_2, client_id="client_id_3", username="admin")
-        self.run_sql(QUERY_3, client_id="client_id_2", username="gamma_sqllab")
+        self.run_sql(QUERY_2, client_id="client_id_2", username="admin")
+        self.run_sql(QUERY_3, client_id="client_id_3", username="gamma_sqllab")
         self.logout()
 
     def tearDown(self):
@@ -406,22 +407,17 @@ class TestSqlLab(SupersetTestCase):
         self.assertEqual(2, len(data))
         self.assertIn("birth", data[0]["sql"])
 
-    def test_search_query_on_time(self):
+    def test_search_query_filter_by_time(self):
         self.run_some_queries()
         self.login("admin")
-        first_query_time = (
-            db.session.query(Query).filter_by(sql=QUERY_1).one()
-        ).start_time
-        second_query_time = (
-            db.session.query(Query).filter_by(sql=QUERY_3).one()
-        ).start_time
-        # Test search queries on time filter
-        from_time = "from={}".format(int(first_query_time))
-        to_time = "to={}".format(int(second_query_time))
-        params = [from_time, to_time]
-        resp = self.get_resp("/superset/search_queries?" + "&".join(params))
-        data = json.loads(resp)
-        self.assertEqual(2, len(data))
+        from_time = floor(
+            (db.session.query(Query).filter_by(sql=QUERY_1).one()).start_time
+        )
+        to_time = ceil(
+            (db.session.query(Query).filter_by(sql=QUERY_2).one()).start_time
+        )
+        url = f"/superset/search_queries?from={from_time}&to={to_time}"
+        assert len(self.client.get(url).json) == 2
 
     def test_search_query_only_owned(self) -> None:
         """