chore: add user agent to Databricks requests

apache · Jul 8, 2022 · f8be567 · f8be567
1 parent 0ce0c6e
commit f8be567
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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