fix: unify environment variables referenced by Databricks docs

Long-term fix will be for Databricks to release a Rust SDK for Unity 😄 Fixes delta-io#1627
rtyler · Sep 27, 2023 · 7abd7d1 · 7abd7d1
1 parent 113fd0f
commit 7abd7d1
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 25 deletions.
diff --git a/python/docs/source/usage.rst b/python/docs/source/usage.rst
@@ -71,15 +71,15 @@ For AWS Glue catalog, use AWS environment variables to authenticate.
     {'id': [5, 7, 9, 5, 6, 7, 8, 9]}
 
 For Databricks Unity Catalog authentication, use environment variables:
-  * DATABRICKS_WORKSPACE_URL (e.g. https://adb-62800498333851.30.azuredatabricks.net)
-  * DATABRICKS_ACCESS_TOKEN
+  * DATABRICKS_HOSAT (e.g. https://adb-62800498333851.30.azuredatabricks.net)
+  * DATABRICKS_TOKEN
 
 .. code-block:: python
     
     >>> import os
     >>> from deltalake import DataCatalog, DeltaTable
-    >>> os.environ['DATABRICKS_WORKSPACE_URL'] = "https://adb-62800498333851.30.azuredatabricks.net"
-    >>> os.environ['DATABRICKS_ACCESS_TOKEN'] = "<DBAT>"
+    >>> os.environ['DATABRICKS_HOST'] = "https://adb-62800498333851.30.azuredatabricks.net"
+    >>> os.environ['DATABRICKS_TOKEN'] = "<DBAT>"
     >>> catalog_name = 'main'
     >>> schema_name = 'db_schema'
     >>> table_name = 'db_table'

diff --git a/rust/src/data_catalog/unity/mod.rs b/rust/src/data_catalog/unity/mod.rs
@@ -81,18 +81,16 @@ pub enum UnityCatalogConfigKey {
     /// Url of a Databricks workspace
     ///
     /// Supported keys:
-    /// - `unity_workspace_url`
-    /// - `databricks_workspace_url`
-    /// - `workspace_url`
-    WorkspaceUrl,
+    /// - `unity_host
+    /// - `databricks_host
+    Host,
 
     /// Access token to authorize API requests
     ///
     /// Supported keys:
-    /// - `unity_access_token`
-    /// - `databricks_access_token`
-    /// - `access_token`
-    AccessToken,
+    /// - `unity_token
+    /// - `databricks_token
+    Token,
 
     /// Service principal client id for authorizing requests
     ///
@@ -169,9 +167,7 @@ impl FromStr for UnityCatalogConfigKey {
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
         match s {
-            "access_token" | "unity_access_token" | "databricks_access_token" => {
-                Ok(UnityCatalogConfigKey::AccessToken)
-            }
+            "unity_token" | "databricks_token" => Ok(UnityCatalogConfigKey::Token),
             "authority_host" | "unity_authority_host" | "databricks_authority_host" => {
                 Ok(UnityCatalogConfigKey::AuthorityHost)
             }
@@ -199,9 +195,7 @@ impl FromStr for UnityCatalogConfigKey {
             "use_azure_cli" | "unity_use_azure_cli" | "databricks_use_azure_cli" => {
                 Ok(UnityCatalogConfigKey::UseAzureCli)
             }
-            "workspace_url" | "unity_workspace_url" | "databricks_workspace_url" => {
-                Ok(UnityCatalogConfigKey::WorkspaceUrl)
-            }
+            "unity_host" | "databricks_host" => Ok(UnityCatalogConfigKey::Host),
             _ => Err(UnityCatalogError::UnknownConfigKey(s.into()).into()),
         }
     }
@@ -210,7 +204,7 @@ impl FromStr for UnityCatalogConfigKey {
 impl AsRef<str> for UnityCatalogConfigKey {
     fn as_ref(&self) -> &str {
         match self {
-            UnityCatalogConfigKey::AccessToken => "unity_access_token",
+            UnityCatalogConfigKey::Token => "unity_token",
             UnityCatalogConfigKey::AuthorityHost => "unity_authority_host",
             UnityCatalogConfigKey::AuthorityId => "unity_authority_id",
             UnityCatalogConfigKey::ClientId => "unity_client_id",
@@ -220,7 +214,7 @@ impl AsRef<str> for UnityCatalogConfigKey {
             UnityCatalogConfigKey::MsiResourceId => "unity_msi_resource_id",
             UnityCatalogConfigKey::ObjectId => "unity_object_id",
             UnityCatalogConfigKey::UseAzureCli => "unity_use_azure_cli",
-            UnityCatalogConfigKey::WorkspaceUrl => "unity_workspace_url",
+            UnityCatalogConfigKey::Host => "unity_host",
         }
     }
 }
@@ -229,7 +223,7 @@ impl AsRef<str> for UnityCatalogConfigKey {
 #[derive(Default)]
 pub struct UnityCatalogBuilder {
     /// Url of a Databricks workspace
-    workspace_url: Option<String>,
+    host: Option<String>,
 
     /// Bearer token
     bearer_token: Option<String>,
@@ -281,8 +275,8 @@ impl UnityCatalogBuilder {
         value: impl Into<String>,
     ) -> DataCatalogResult<Self> {
         match UnityCatalogConfigKey::from_str(key.as_ref())? {
-            UnityCatalogConfigKey::WorkspaceUrl => self.workspace_url = Some(value.into()),
-            UnityCatalogConfigKey::AccessToken => self.bearer_token = Some(value.into()),
+            UnityCatalogConfigKey::Host => self.host = Some(value.into()),
+            UnityCatalogConfigKey::Token => self.bearer_token = Some(value.into()),
             UnityCatalogConfigKey::ClientId => self.client_id = Some(value.into()),
             UnityCatalogConfigKey::ClientSecret => self.client_secret = Some(value.into()),
             UnityCatalogConfigKey::AuthorityId => self.authority_id = Some(value.into()),
@@ -331,7 +325,7 @@ impl UnityCatalogBuilder {
 
     /// Set the URL of a Databricks workspace.
     pub fn with_workspace_url(mut self, url: impl Into<String>) -> Self {
-        self.workspace_url = Some(url.into());
+        self.host = Some(url.into());
         self
     }
 
@@ -413,7 +407,7 @@ impl UnityCatalogBuilder {
             .ok_or(UnityCatalogError::MissingCredential)?;
 
         let workspace_url = self
-            .workspace_url
+            .host
             .ok_or(UnityCatalogError::MissingConfiguration(
                 "workspace_url".into(),
             ))?