From f8be567f2379b0b55962827ab5d037ba5b6ea9c4 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Fri, 8 Jul 2022 14:40:19 -0700
Subject: [PATCH] chore: add user agent to Databricks requests

---
 docs/docs/databases/databricks.mdx     |  5 +----
 superset/constants.py                  |  2 ++
 superset/db_engine_specs/databricks.py | 18 +++++++++++++++++-
 3 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/docs/docs/databases/databricks.mdx b/docs/docs/databases/databricks.mdx
index 4070960ce2af5..e9bbd4b685219 100644
--- a/docs/docs/databases/databricks.mdx
+++ b/docs/docs/databases/databricks.mdx
@@ -33,13 +33,10 @@ You also need to add the following configuration to "Other" -> "Engine Parameter
 
 ```json
 {
-    "connect_args": {"http_path": "sql/protocolv1/o/****"},
-    "http_headers": [["User-Agent", "Apache Superset"]]
+    "connect_args": {"http_path": "sql/protocolv1/o/****"}
 }
 ```
 
-The `User-Agent` header is optional, but helps Databricks identify traffic from Superset. If you need to use a different header please reach out to Databricks and let them know.
-
 ## Older driver
 
 Originally Superset used `databricks-dbapi` to connect to Databricks. You might want to try it if you're having problems with the official Databricks connector:
diff --git a/superset/constants.py b/superset/constants.py
index 5df4f961e6859..32e0a9abce3ca 100644
--- a/superset/constants.py
+++ b/superset/constants.py
@@ -20,6 +20,8 @@
 # string to use when None values *need* to be converted to/from strings
 from enum import Enum
 
+USER_AGENT = "Apache Superset"
+
 NULL_STRING = "<NULL>"
 EMPTY_STRING = "<empty string>"
 
diff --git a/superset/db_engine_specs/databricks.py b/superset/db_engine_specs/databricks.py
index 8803d0522e2ed..9bb00857544c1 100644
--- a/superset/db_engine_specs/databricks.py
+++ b/superset/db_engine_specs/databricks.py
@@ -16,11 +16,15 @@
 # under the License.
 
 from datetime import datetime
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, TYPE_CHECKING
 
+from superset.constants import USER_AGENT
 from superset.db_engine_specs.base import BaseEngineSpec
 from superset.db_engine_specs.hive import HiveEngineSpec
 
+if TYPE_CHECKING:
+    from superset.models.core import Database
+
 time_grain_expressions = {
     None: "{col}",
     "PT1S": "date_trunc('second', {col})",
@@ -71,3 +75,15 @@ class DatabricksNativeEngineSpec(DatabricksODBCEngineSpec):
     engine = "databricks"
     engine_name = "Databricks Native Connector"
     driver = "connector"
+
+    @staticmethod
+    def get_extra_params(database: "Database") -> Dict[str, Any]:
+        """
+        Add a user agent to be used in the requests.
+        """
+        extra = {
+            "http_headers": [("User-Agent", USER_AGENT)],
+            "_user_agent_entry": USER_AGENT,
+        }
+        extra.update(BaseEngineSpec.get_extra_params(database))
+        return extra