Fix #2984: added azure sso auth (#3002)

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>"
+        ]
+      }
+    }
+  }

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",

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)

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)

ingestion/src/metadata/ingestion/sink/metadata_rest.py

@@ -419,6 +419,7 @@ class MetadataRestSink(Sink[Entity]):
                     )
                 except APIError:
                     role_entity = self._create_role(role)
+                if role_entity:
                     role_ids.append(role_entity.id)
         else:
             role_ids = None