From 0b44bbdde1e57670638329dafdba674a6746725d Mon Sep 17 00:00:00 2001
From: Ayush Shah <ayush@getcollate.io>
Date: Sat, 9 Apr 2022 07:58:40 -0700
Subject: [PATCH] re: Azure SSO added (#3989)

* Azure SSO added

* Azure SSO added
---
 .../schema/metadataIngestion/workflow.json    | 29 +++++++++++++++++++
 .../metadata/ingestion/ometa/auth_provider.py | 11 +++----
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/catalog-rest-service/src/main/resources/json/schema/metadataIngestion/workflow.json b/catalog-rest-service/src/main/resources/json/schema/metadataIngestion/workflow.json
index 45f17ded454..822e9b23c11 100644
--- a/catalog-rest-service/src/main/resources/json/schema/metadataIngestion/workflow.json
+++ b/catalog-rest-service/src/main/resources/json/schema/metadataIngestion/workflow.json
@@ -5,6 +5,32 @@
   "description": "OpenMetadata Ingestion Framework definition.",
   "type": "object",
   "definitions": {
+    "azureSSOConfig": {
+      "description": "Azure SSO client security configs.",
+      "type": "object",
+      "properties": {
+        "clientSecret": {
+          "description": "Azure SSO client secret key",
+          "type": "string"
+        },
+        "authority": {
+          "description": "Azure SSO Authority",
+          "type": "string"
+        },
+        "clientId": {
+          "description": "Azure Client ID.",
+          "type": "string"
+        },
+        "scopes": {
+          "description": "Azure Client ID.",
+          "items": {
+            "type": "string"
+          }
+        }
+      },
+      "additionalProperties": false,
+      "required": ["clientSecret", "clientId", "authority", "scopes"]
+    },
     "googleSSOConfig": {
       "description": "Google SSO client security configs.",
       "type": "object",
@@ -120,6 +146,9 @@
             {
               "$ref": "#/definitions/auth0SSOConfig"
             },
+            {
+              "$ref": "#/definitions/azureSSOConfig"
+            },
             {
               "$ref": "#/definitions/customOidcSSOConfig"
             }
diff --git a/ingestion/src/metadata/ingestion/ometa/auth_provider.py b/ingestion/src/metadata/ingestion/ometa/auth_provider.py
index e662934bfb5..ec9c19ea017 100644
--- a/ingestion/src/metadata/ingestion/ometa/auth_provider.py
+++ b/ingestion/src/metadata/ingestion/ometa/auth_provider.py
@@ -25,6 +25,7 @@ import requests
 from metadata.config.common import ConfigModel
 from metadata.generated.schema.metadataIngestion.workflow import (
     Auth0SSOConfig,
+    AzureSSOConfig,
     CustomOidcSSOConfig,
     GoogleSSOConfig,
     OktaSSOConfig,
@@ -278,7 +279,7 @@ class AzureAuthenticationProvider(AuthenticationProvider):
     # TODO: Prepare JSON for Azure Auth
     def __init__(self, config: OpenMetadataServerConfig):
         self.config = config
-
+        self.security_config: AzureSSOConfig = self.config.securityConfig
         self.generated_auth_token = None
         self.expiry = None
 
@@ -292,11 +293,11 @@ class AzureAuthenticationProvider(AuthenticationProvider):
         )
 
         app = ConfidentialClientApplication(
-            client_id=self.config.client_id,
-            client_credential=self.config.secret_key,
-            authority=self.config.authority,
+            client_id=self.security_config.clientId,
+            client_credential=self.security_config.clientSecret,
+            authority=self.security_config.authority,
         )
-        token = app.acquire_token_for_client(scopes=self.config.scopes)
+        token = app.acquire_token_for_client(scopes=self.security_config.scopes)
         try:
             self.generated_auth_token = token["access_token"]
             self.expiry = token["expires_in"]