From 68bd1b542d7221e863a0083e67e1bca31cf445a4 Mon Sep 17 00:00:00 2001
From: Serge Smertin <serge.smertin@databricks.com>
Date: Wed, 17 May 2023 11:30:50 +0200
Subject: [PATCH 1/3] Prototype for notebook-native auth

Fixes #57
---
 databricks/sdk/core.py             | 13 ++++++++++++-
 databricks/sdk/runtime/__init__.py |  8 ++++++++
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/databricks/sdk/core.py b/databricks/sdk/core.py
index f8bd89eb0..e9c776d59 100644
--- a/databricks/sdk/core.py
+++ b/databricks/sdk/core.py
@@ -87,6 +87,17 @@ def inner() -> Dict[str, str]:
     return inner
 
 
+@credentials_provider('runtime', [])
+def runtime_native_auth(cfg: 'Config') -> Optional[HeaderFactory]:
+    from databricks.sdk.runtime import init_runtime_native_auth
+    try:
+        host, inner = init_runtime_native_auth()
+        cfg.host = host
+        return inner
+    except NotImplemented:
+        return None
+
+
 @credentials_provider('oauth-m2m', ['is_aws', 'host', 'client_id', 'client_secret'])
 def oauth_service_principal(cfg: 'Config') -> Optional[HeaderFactory]:
     """ Adds refreshed Databricks machine-to-machine OAuth Bearer token to every request,
@@ -309,7 +320,7 @@ def auth_type(self) -> str:
     def __call__(self, cfg: 'Config') -> HeaderFactory:
         auth_providers = [
             pat_auth, basic_auth, oauth_service_principal, azure_service_principal, azure_cli,
-            external_browser, bricks_cli
+            external_browser, bricks_cli, runtime_native_auth
         ]
         for provider in auth_providers:
             auth_type = provider.auth_type()
diff --git a/databricks/sdk/runtime/__init__.py b/databricks/sdk/runtime/__init__.py
index 2bfcc4939..e666be1fe 100644
--- a/databricks/sdk/runtime/__init__.py
+++ b/databricks/sdk/runtime/__init__.py
@@ -7,6 +7,14 @@
     "display", "displayHTML", "dbutils", "table", "sql", "udf", "getArgument", "sc", "sqlContext", "spark"
 ]
 
+from typing import Dict, Tuple, Callable
+
+RuntimeAuth = Tuple[str, Callable[[], Dict[str, str]]]
+
+
+def init_runtime_native_auth() -> RuntimeAuth:
+    raise NotImplemented
+
 try:
     # Internal implementation
     from dbruntime import UserNamespaceInitializer

From 1b0f77a8a5c917fa7b503a5a2569b5d70feca888 Mon Sep 17 00:00:00 2001
From: Serge Smertin <serge.smertin@databricks.com>
Date: Wed, 17 May 2023 11:50:43 +0200
Subject: [PATCH 2/3] make fmt

---
 databricks/sdk/core.py             | 6 +++++-
 databricks/sdk/runtime/__init__.py | 3 ++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/databricks/sdk/core.py b/databricks/sdk/core.py
index e9c776d59..682aff6c9 100644
--- a/databricks/sdk/core.py
+++ b/databricks/sdk/core.py
@@ -765,7 +765,11 @@ def _authenticate(self, r: requests.PreparedRequest) -> requests.PreparedRequest
 
     def do(self, method: str, path: str, query: dict = None, body: dict = None) -> dict:
         headers = {'Accept': 'application/json', 'User-Agent': self._user_agent_base}
-        response = self._session.request(method, f"{self._cfg.host}{path}", params=query, json=body, headers=headers)
+        response = self._session.request(method,
+                                         f"{self._cfg.host}{path}",
+                                         params=query,
+                                         json=body,
+                                         headers=headers)
         try:
             self._record_request_log(response)
             if not response.ok:
diff --git a/databricks/sdk/runtime/__init__.py b/databricks/sdk/runtime/__init__.py
index e666be1fe..1f8bbd0d6 100644
--- a/databricks/sdk/runtime/__init__.py
+++ b/databricks/sdk/runtime/__init__.py
@@ -7,7 +7,7 @@
     "display", "displayHTML", "dbutils", "table", "sql", "udf", "getArgument", "sc", "sqlContext", "spark"
 ]
 
-from typing import Dict, Tuple, Callable
+from typing import Callable, Dict, Tuple
 
 RuntimeAuth = Tuple[str, Callable[[], Dict[str, str]]]
 
@@ -15,6 +15,7 @@
 def init_runtime_native_auth() -> RuntimeAuth:
     raise NotImplemented
 
+
 try:
     # Internal implementation
     from dbruntime import UserNamespaceInitializer

From b94e632aea419d4d9271e1538491a8acc3b6aaa3 Mon Sep 17 00:00:00 2001
From: Serge Smertin <serge.smertin@databricks.com>
Date: Wed, 17 May 2023 11:53:05 +0200
Subject: [PATCH 3/3] fix exception

---
 databricks/sdk/core.py             | 2 +-
 databricks/sdk/runtime/__init__.py | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/databricks/sdk/core.py b/databricks/sdk/core.py
index 682aff6c9..f72addd9f 100644
--- a/databricks/sdk/core.py
+++ b/databricks/sdk/core.py
@@ -94,7 +94,7 @@ def runtime_native_auth(cfg: 'Config') -> Optional[HeaderFactory]:
         host, inner = init_runtime_native_auth()
         cfg.host = host
         return inner
-    except NotImplemented:
+    except NotImplementedError:
         return None
 
 
diff --git a/databricks/sdk/runtime/__init__.py b/databricks/sdk/runtime/__init__.py
index 1f8bbd0d6..e83c9d65c 100644
--- a/databricks/sdk/runtime/__init__.py
+++ b/databricks/sdk/runtime/__init__.py
@@ -1,3 +1,5 @@
+from typing import Callable, Dict, Tuple
+
 is_local_implementation = True
 
 # All objects that are injected into the Notebook's user namespace should also be made
@@ -7,13 +9,11 @@
     "display", "displayHTML", "dbutils", "table", "sql", "udf", "getArgument", "sc", "sqlContext", "spark"
 ]
 
-from typing import Callable, Dict, Tuple
-
 RuntimeAuth = Tuple[str, Callable[[], Dict[str, str]]]
 
 
 def init_runtime_native_auth() -> RuntimeAuth:
-    raise NotImplemented
+    raise NotImplementedError
 
 
 try: