Fixed Interactive OAuth on Azure & updated documentations (#669)

nkvuong · web-flow · commit 455a14c1e345 · 2024-06-07T09:44:21.000Z
## Changes
- Removed Azure check in interactive OAuth
- Updated Flask with OAuth example
- Updated OAuth documentation

## Tests
&lt;!-- 
How is this tested? Please see the checklist below and also describe any
other relevant tests
--&gt;

- [x] `make test` run locally
- [x] `make fmt` applied
- [ ] relevant integration tests applied
diff --git a/databricks/sdk/oauth.py b/databricks/sdk/oauth.py
@@ -371,9 +371,6 @@ def noop_credentials(_: any):
         config = Config(host=host, credentials_strategy=noop_credentials)
         if not scopes:
             scopes = ['all-apis']
-        if config.is_azure:
-            # Azure AD only supports full access to Azure Databricks.
-            scopes = [f'{config.effective_azure_login_app_id}/user_impersonation', 'offline_access']
         oidc = config.oidc_endpoints
         if not oidc:
             raise ValueError(f'{host} does not support OAuth')
@@ -385,6 +382,7 @@ def noop_credentials(_: any):
         self.token_url = oidc.token_endpoint
         self.is_aws = config.is_aws
         self.is_azure = config.is_azure
+        self.is_gcp = config.is_gcp
 
         self._auth_url = oidc.authorization_endpoint
         self._scopes = scopes
diff --git a/docs/oauth.md b/docs/oauth.md
@@ -245,8 +245,6 @@ Databricks SDK for Python exposes the `oauth_client.initiate_consent()` helper t
 PKCE state verification. Application developers are expected to persist `RefreshableCredentials` in the webapp session
 and restore it via `RefreshableCredentials.from_dict(oauth_client, session['creds'])` helpers.
 
-Works for both AWS and Azure. Not supported for GCP at the moment.
-
 ```python
 from databricks.sdk.oauth import OAuthClient
 oauth_client = OAuthClient(host='<workspace-url>',
@@ -305,10 +303,6 @@ account_client = AccountClient(host='https://accounts.cloud.databricks.com',
                                account_id=input('Databricks Account ID: '),
                                username=input('Username: '),
                                password=getpass.getpass('Password: '))
-logging.info('Enrolling all published apps...')
-account_client.o_auth_enrollment.create(enable_all_published_apps=True)
-status = account_client.o_auth_enrollment.get()
-logging.info(f'Enrolled all published apps: {status}')
 custom_app = account_client.custom_app_integration.create(
     name='awesome-app',
     redirect_urls=[f'https://host.domain/path/to/callback'],
diff --git a/examples/flask_app_with_oauth.py b/examples/flask_app_with_oauth.py
@@ -8,14 +8,14 @@
 
 If you have already Custom App:
 
-./flask_app_with_oauth.py <databricks workspace url> \
+./flask_app_with_oauth.py --host <databricks workspace url> \
     --client_id <app-client-id> \
     --client_secret <app-secret> \
     --port 5001
 
 If you want this script to register Custom App and redirect URL for you:
 
-./flask_app_with_oauth.py <databricks workspace url>
+./flask_app_with_oauth.py --port 5001 --profile <databricks account profile>
 
 You'll get prompted for Databricks Account username and password for
 script to enroll your account into OAuth and create a custom app with
@@ -44,7 +44,7 @@
 </ul>"""
 
 
-def create_flask_app(oauth_client: OAuthClient, port: int):
+def create_flask_app(oauth_client: OAuthClient):
     """The create_flask_app function creates a Flask app that is enabled with OAuth.
 
     It initializes the app and web session secret keys with a randomly generated token. It defines two routes for
@@ -91,23 +91,13 @@ def index():
     return app
 
 
-def register_custom_app(oauth_client: OAuthClient, args: argparse.Namespace) -> tuple[str, str]:
+def register_custom_app(args: argparse.Namespace) -> tuple[str, str]:
     """Creates new Custom OAuth App in Databricks Account"""
-    if not oauth_client.is_aws:
-        logging.error("Not supported for other clouds than AWS")
-        sys.exit(2)
-
     logging.info("No OAuth custom app client/secret provided, creating new app")
 
-    import getpass
-
     from databricks.sdk import AccountClient
 
-    account_client = AccountClient(host="https://accounts.cloud.databricks.com",
-                                   account_id=input("Databricks Account ID: "),
-                                   username=input("Username: "),
-                                   password=getpass.getpass("Password: "),
-                                   )
+    account_client = AccountClient(profile=args.profile)
 
     custom_app = account_client.custom_app_integration.create(
         name=APP_NAME, redirect_urls=[f"http://localhost:{args.port}/callback"], confidential=True,
@@ -129,7 +119,7 @@ def init_oauth_config(args) -> OAuthClient:
                                scopes=["all-apis"],
                                )
     if not oauth_client.client_id:
-        client_id, client_secret = register_custom_app(oauth_client, args)
+        client_id, client_secret = register_custom_app(args)
         oauth_client.client_id = client_id
         oauth_client.client_secret = client_secret
 
@@ -139,10 +129,11 @@ def init_oauth_config(args) -> OAuthClient:
 def parse_arguments() -> argparse.Namespace:
     """Parses arguments for this demo"""
     parser = argparse.ArgumentParser(prog=APP_NAME, description=__doc__.strip())
-    parser.add_argument("host")
+    parser.add_argument("--host")
     for flag in ["client_id", "client_secret"]:
         parser.add_argument(f"--{flag}")
     parser.add_argument("--port", default=5001, type=int)
+    parser.add_argument("--profile", default="DEFAULT", help="Databricks account profile to use for authentication.")
     return parser.parse_args()
 
 
@@ -155,7 +146,7 @@ def parse_arguments() -> argparse.Namespace:
 
     args = parse_arguments()
     oauth_cfg = init_oauth_config(args)
-    app = create_flask_app(oauth_cfg, args.port)
+    app = create_flask_app(oauth_cfg)
 
     app.run(
         host="localhost",
