Merge 3b9dd7b347 into f5843fe588

superset/db_engine_specs/snowflake.py

@@ -16,6 +16,7 @@
 # under the License.
 import json
 import logging
+import os
 import re
 from datetime import datetime
 from re import Pattern
@@ -274,6 +275,7 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
                 query={
                     "role": parameters.get("role"),
                     "warehouse": parameters.get("warehouse"),
+                    "application": os.environ.get("SF_PARTNER", "SUPERSET"),
                 },
             )
         )

tests/unit_tests/db_engine_specs/test_snowflake.py

@@ -147,6 +147,30 @@ def test_cancel_query_failed(engine_mock: mock.Mock) -> None:
     assert SnowflakeEngineSpec.cancel_query(cursor_mock, query, "123") is False
 
 
+def test_build_sqlalchemy_uri() -> None:
+    from superset.db_engine_specs.snowflake import (
+        SnowflakeEngineSpec,
+        SnowflakeParametersType,
+    )
+
+    parameters: SnowflakeParametersType = {
+        "username": "test-username",
+        "password": "test-password",
+        "account": "test-account",
+        "database": "test-database",
+        "role": "test-role",
+        "warehouse": "test-warehouse",
+    }
+    assert SnowflakeEngineSpec.build_sqlalchemy_uri(parameters) == (
+        "snowflake://test-username:test-password@test-account/test-database?"
+        "role=test-role&warehouse=test-warehouse&application=SUPERSET"
+    )
+    with mock.patch.dict("os.environ", SF_PARTNER="PARTNER"):
+        assert SnowflakeEngineSpec.build_sqlalchemy_uri(parameters) == (
+            "snowflake://test-username:test-password@test-account/test-database?"
+            "role=test-role&warehouse=test-warehouse&application=PARTNER"
+        )
+
 def test_get_extra_params(mocker: MockerFixture) -> None:
     """
     Test the ``get_extra_params`` method.