Improved notebook-native authentication #152
Conversation
There was a problem hiding this comment.
Can we move all the init_runtime_native_auth and databricks_repl_context logic to a new file, there is no reason any of this needs to be databricks.sdk.runtime. The auth stuff and the runtime stuff are distinct and we explicitly do not want them to interact. All this code should work the same in another file.
databricks/sdk/runtime/__init__.py
Outdated
| try: | ||
| from dbruntime.databricks_repl_context import get_context | ||
| ctx = get_context() | ||
| if ctx is not None: |
There was a problem hiding this comment.
This is a static check on ctx done on import, but get_context will not return a Context until a command has been run. This probably works because databricks.sdk.rutnime is probably not imported until a user interacts with the sdk, but this is pretty brittle and will break if we internally import this file on startup in the future.
Can we make this check either dynamic or at if we want to check once and then give up, could we be more explcit about when the check is done.
There was a problem hiding this comment.
addressed in f371865 - we need a hostname during the SDK initialization stage. technically, if we can start every runtime with DATABRICKS_HOST environment variable with the f'https://{ctx.browserHostName}', we can simplify the init_runtime_native_auth return type.
please remember - init_runtime_native_auth returns the workspace hostname and function that returns a dictionary of headers. this is to support future improvements to REPL authentication, where we could leverage token refreshes
…se to other credentials providers
Codecov ReportPatch coverage:
Additional details and impacted files@@ Coverage Diff @@
## main #152 +/- ##
==========================================
+ Coverage 53.23% 53.25% +0.01%
==========================================
Files 29 30 +1
Lines 17917 18190 +273
==========================================
+ Hits 9539 9687 +148
- Misses 8378 8503 +125
☔ View full report in Codecov by Sentry. |
| host, inner = init_runtime_native_auth() | ||
| cfg.host = host | ||
| return inner | ||
| try: | ||
| from dbruntime.databricks_repl_context import get_context |
There was a problem hiding this comment.
Hi @nfx may I know whether you verify in notebook environment (e.g. with DBR 12.2) that it works? I vaguely remember last time I want to import from dbruntime in core.py it does not work. Ignore me if I'm wrong.
There was a problem hiding this comment.
I've tested it on couple of DBRs
| cfg.host = f'https://{ctx.browserHostName}' | ||
|
|
||
| def inner() -> Dict[str, str]: | ||
| return {'Authorization': f'Bearer {ctx.apiToken}'} |
There was a problem hiding this comment.
Don't we lock in the apiToken? I think we should call get_context() inside inner to swap the context for different users.
There was a problem hiding this comment.
@xiaochen-db also running get_context() within the inner? thanks for pointing this out! will add
| host, inner = init_runtime_native_auth() | ||
| cfg.host = host | ||
| return inner | ||
| try: | ||
| from dbruntime.databricks_repl_context import get_context | ||
| ctx = get_context() |
There was a problem hiding this comment.
Just want to double check when is this called? Inside a command or during REPL initialization?
* Regenerate from OpenAPI spec ([#176](#176)). * Added improved notebook-native authentication ([#152](#152)). * Added methods to provide extra user agent and upstream user agent to SDK config ([#163](#163)). * Added more missing `Optional` type hints ([#171](#171)). * Add more missing optional fields ([#177](#177)). * Correctly serialize external entities ([#178](#178)). * Correctly serialize external enum values in paths ([#179](#179)). * Mark non-required fields as `Optional` ([#170](#170)). * Synchronize auth permutation tests with Go SDK ([#165](#165)).
* Regenerate from OpenAPI spec ([#176](#176)). * Added improved notebook-native authentication ([#152](#152)). * Added methods to provide extra user agent and upstream user agent to SDK config ([#163](#163)). * Added more missing `Optional` type hints ([#171](#171)). * Added more missing optional fields ([#177](#177)). * Correctly serialize external entities ([#178](#178)). * Correctly serialize external enum values in paths ([#179](#179)). * Mark non-required fields as `Optional` ([#170](#170)). * Synchronize auth permutation tests with Go SDK ([#165](#165)). ## Tests - [x] relevant integration tests applied --------- Signed-off-by: Serge Smertin <259697+nfx@users.noreply.github.com>
Ported from MLflow codebase:
Resolves #153.