From c49af971a75f6eca5b940c3ea078f7dc028cfe60 Mon Sep 17 00:00:00 2001
From: Mayur Singal <39544459+ulixius9@users.noreply.github.com>
Date: Tue, 1 Mar 2022 01:14:28 +0530
Subject: [PATCH] Fix #2984: added azure sso auth (#3002)

---
 ingestion/examples/auth_examples/azure.json   | 25 +++++++++++
 ingestion/setup.py                            |  1 +
 .../src/metadata/ingestion/ometa/ometa_api.py |  5 +++
 .../ingestion/ometa/openmetadata_rest.py      | 41 +++++++++++++++++++
 .../metadata/ingestion/sink/metadata_rest.py  |  3 +-
 5 files changed, 74 insertions(+), 1 deletion(-)
 create mode 100644 ingestion/examples/auth_examples/azure.json

diff --git a/ingestion/examples/auth_examples/azure.json b/ingestion/examples/auth_examples/azure.json
new file mode 100644
index 00000000000..cd04591fe1e
--- /dev/null
+++ b/ingestion/examples/auth_examples/azure.json
@@ -0,0 +1,25 @@
+{
+    "source": {
+      "type": "sample-data",
+      "config": {
+        "sample_data_folder": "./examples/sample_data"
+      }
+    },
+    "sink": {
+      "type": "metadata-rest",
+      "config": {}
+    },
+    "metadata_server": {
+      "type": "metadata-server",
+      "config": {
+        "api_endpoint": "http://localhost:8585/api",
+        "auth_provider_type": "azure",
+        "client_id": "<client_id>",
+        "authority":"https://login.microsoftonline.com/<tenant_id>",
+        "secret_key":"<client_secret>",
+        "scopes": [
+            "<resource_uri>"
+        ]
+      }
+    }
+  }
\ No newline at end of file
diff --git a/ingestion/setup.py b/ingestion/setup.py
index bf45e4ab2c5..14bc0d2b6ff 100644
--- a/ingestion/setup.py
+++ b/ingestion/setup.py
@@ -124,6 +124,7 @@ plugins: Dict[str, Set[str]] = {
     "clickhouse": {"clickhouse-driver==0.2.3", "clickhouse-sqlalchemy==0.2.0"},
     "databricks": {"sqlalchemy-databricks==0.1.0"},
     "singlestore": {"pymysql>=1.0.2"},
+    "azure-sso": {"msal~=1.17.0"},
 }
 dev = {
     "boto3==1.20.14",
diff --git a/ingestion/src/metadata/ingestion/ometa/ometa_api.py b/ingestion/src/metadata/ingestion/ometa/ometa_api.py
index 0b73fe27cbd..902e4e3c9b0 100644
--- a/ingestion/src/metadata/ingestion/ometa/ometa_api.py
+++ b/ingestion/src/metadata/ingestion/ometa/ometa_api.py
@@ -57,6 +57,7 @@ from metadata.ingestion.ometa.mixins.tag_mixin import OMetaTagMixin
 from metadata.ingestion.ometa.mixins.version_mixin import OMetaVersionMixin
 from metadata.ingestion.ometa.openmetadata_rest import (
     Auth0AuthenticationProvider,
+    AzureAuthenticationProvider,
     GoogleAuthenticationProvider,
     MetadataServerConfig,
     NoOpAuthenticationProvider,
@@ -148,6 +149,10 @@ class OpenMetadata(
             self._auth_provider: AuthenticationProvider = (
                 Auth0AuthenticationProvider.create(self.config)
             )
+        elif self.config.auth_provider_type == "azure":
+            self._auth_provider: AuthenticationProvider = (
+                AzureAuthenticationProvider.create(self.config)
+            )
         else:
             self._auth_provider: AuthenticationProvider = (
                 NoOpAuthenticationProvider.create(self.config)
diff --git a/ingestion/src/metadata/ingestion/ometa/openmetadata_rest.py b/ingestion/src/metadata/ingestion/ometa/openmetadata_rest.py
index 4473228ae3c..143bddf7813 100644
--- a/ingestion/src/metadata/ingestion/ometa/openmetadata_rest.py
+++ b/ingestion/src/metadata/ingestion/ometa/openmetadata_rest.py
@@ -16,7 +16,9 @@ import http.client
 import json
 import logging
 import sys
+import time
 import traceback
+import uuid
 from typing import List
 
 from pydantic import BaseModel
@@ -98,6 +100,7 @@ class MetadataServerConfig(ConfigModel):
     email: str = None
     audience: str = "https://www.googleapis.com/oauth2/v4/token"
     auth_header: str = "Authorization"
+    authority: str = ""
     scopes: List = []
 
 
@@ -282,3 +285,41 @@ class Auth0AuthenticationProvider(AuthenticationProvider):
     def get_access_token(self):
         self.auth_token()
         return (self.generated_auth_token, self.expiry)
+
+
+class AzureAuthenticationProvider(AuthenticationProvider):
+    """
+    Prepare the Json Web Token for Azure auth
+    """
+
+    def __init__(self, config: MetadataServerConfig):
+        self.config = config
+
+    @classmethod
+    def create(cls, config: MetadataServerConfig):
+        return cls(config)
+
+    def auth_token(self) -> str:
+        from msal import (
+            ConfidentialClientApplication,  # pylint: disable=import-outside-toplevel
+        )
+
+        app = ConfidentialClientApplication(
+            client_id=self.config.client_id,
+            client_credential=self.config.secret_key,
+            authority=self.config.authority,
+        )
+        token = app.acquire_token_for_client(scopes=self.config.scopes)
+        try:
+            self.generated_auth_token = token["access_token"]
+            self.expiry = token["expires_in"]
+
+        except KeyError as err:
+            logger.error(f"Invalid Credentials - {err}")
+            logger.debug(traceback.format_exc())
+            logger.debug(traceback.print_exc())
+            sys.exit(1)
+
+    def get_access_token(self):
+        self.auth_token()
+        return (self.generated_auth_token, self.expiry)
diff --git a/ingestion/src/metadata/ingestion/sink/metadata_rest.py b/ingestion/src/metadata/ingestion/sink/metadata_rest.py
index 907ebe6d9ff..cb166b3b837 100644
--- a/ingestion/src/metadata/ingestion/sink/metadata_rest.py
+++ b/ingestion/src/metadata/ingestion/sink/metadata_rest.py
@@ -419,7 +419,8 @@ class MetadataRestSink(Sink[Entity]):
                     )
                 except APIError:
                     role_entity = self._create_role(role)
-                role_ids.append(role_entity.id)
+                if role_entity:
+                    role_ids.append(role_entity.id)
         else:
             role_ids = None