From 69c45ff835b08a83619109edde590d88b3cff809 Mon Sep 17 00:00:00 2001
From: David Blain <david.blain@infrabel.be>
Date: Tue, 17 Dec 2024 20:08:49 +0100
Subject: [PATCH] refactor: Added MS Graph connection type

---
 .../microsoft/azure/hooks/msgraph.py          | 48 +++++++++++++++++++
 .../providers/microsoft/azure/provider.yaml   |  2 +
 2 files changed, 50 insertions(+)

diff --git a/providers/src/airflow/providers/microsoft/azure/hooks/msgraph.py b/providers/src/airflow/providers/microsoft/azure/hooks/msgraph.py
index 4ab3aaf3ba37f..9f97f862089f2 100644
--- a/providers/src/airflow/providers/microsoft/azure/hooks/msgraph.py
+++ b/providers/src/airflow/providers/microsoft/azure/hooks/msgraph.py
@@ -108,7 +108,10 @@ class KiotaRequestAdapterHook(BaseHook):
 
     DEFAULT_HEADERS = {"Accept": "application/json;q=1"}
     cached_request_adapters: dict[str, tuple[APIVersion, RequestAdapter]] = {}
+    conn_type: str = "msgraph"
+    conn_name_attr: str = "conn_id"
     default_conn_name: str = "msgraph_default"
+    hook_name: str = "Microsoft Graph API"
 
     def __init__(
         self,
@@ -127,6 +130,51 @@ def __init__(
         self.scopes = scopes or ["https://graph.microsoft.com/.default"]
         self._api_version = self.resolve_api_version_from_value(api_version)
 
+    @classmethod
+    def get_connection_form_widgets(cls) -> dict[str, Any]:
+        """Return connection widgets to add to connection form."""
+        from flask_appbuilder.fieldwidgets import BS3TextAreaFieldWidget, BS3TextFieldWidget
+        from flask_babel import lazy_gettext
+        from wtforms.fields import BooleanField, StringField
+
+        return {
+            "tenant_id": StringField(lazy_gettext("Tenant ID"), widget=BS3TextFieldWidget()),
+            "api_version": StringField(
+                lazy_gettext("API Version"), widget=BS3TextFieldWidget(), default="v1.0"
+            ),
+            "authority": StringField(lazy_gettext("Authority"), widget=BS3TextFieldWidget()),
+            "scopes": StringField(
+                lazy_gettext("Scopes"),
+                widget=BS3TextFieldWidget(),
+                default="https://graph.microsoft.com/.default",
+            ),
+            "disable_instance_discovery": BooleanField(
+                lazy_gettext("Disable instance discovery"), default=False
+            ),
+            "allowed_hosts": StringField(lazy_gettext("Allowed"), widget=BS3TextFieldWidget()),
+            "proxies": StringField(lazy_gettext("Proxies"), widget=BS3TextAreaFieldWidget()),
+            "stream": BooleanField(lazy_gettext("Stream"), default=False),
+            "verify": BooleanField(lazy_gettext("Verify"), default=True),
+            "trust_env": BooleanField(lazy_gettext("Trust environment"), default=True),
+            "base_url": StringField(lazy_gettext("Base URL"), widget=BS3TextFieldWidget()),
+        }
+
+    @classmethod
+    def get_ui_field_behaviour(cls) -> dict[str, Any]:
+        """Return custom field behaviour."""
+        return {
+            "hidden_fields": ["extra"],
+            "relabeling": {
+                "login": "Client ID",
+                "password": "Client Secret",
+            },
+            "default_values": {
+                "schema": "https",
+                "host": NationalClouds.Global.value,
+                "port": 443,
+            },
+        }
+
     @property
     def api_version(self) -> str | None:
         self.get_conn()  # Make sure config has been loaded through get_conn to have correct api version!
diff --git a/providers/src/airflow/providers/microsoft/azure/provider.yaml b/providers/src/airflow/providers/microsoft/azure/provider.yaml
index 8db33d42ff3ba..5f42247ee2011 100644
--- a/providers/src/airflow/providers/microsoft/azure/provider.yaml
+++ b/providers/src/airflow/providers/microsoft/azure/provider.yaml
@@ -359,6 +359,8 @@ connection-types:
     connection-type: azure_synapse
   - hook-class-name: airflow.providers.microsoft.azure.hooks.data_lake.AzureDataLakeStorageV2Hook
     connection-type: adls
+  - hook-class-name: airflow.providers.microsoft.azure.hooks.msgraph.KiotaRequestAdapterHook
+    connection-type: msgraph
   - hook-class-name: airflow.providers.microsoft.azure.hooks.powerbi.PowerBIHook
     connection-type: powerbi