Databricks のデータフレームにて Azure Event Hubs から OAuth 認証によりデータを取得する方法

## 概要

Databricks のデータフレームにて Azure Event Hubs から OAuth 認証によりデータを取得する方法を紹介します。

## 前提条件

下記を構築します。

- Databricks
- Azure Event Hubs 名前空間（本手順では`kafka-test-01`という名称）
- Azure Event Hubs インスタンス （本手順では`dbx-to-eh-01`という名称）

Azure Event Hubs インスタンス の Data Explorer にて、イエロー タクシー データセットのサンプルデータを送信します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/ce3d7529-5610-4343-866a-380c1657473e.png)

## 手順

### 1. Azure Entra ID にてアプリを作成

Azure Poral にて Azure Entra ID を表示し、`アプリの登録` -> `+ 新規登録`を選択します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/adeccf1c-1837-4205-b46b-f8ea4336e4d1.png)

名前欄を入力し、`登録`を選択します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/da5de324-3b0d-4581-8ea0-60572e615c52.png)

作成したアプリにおける`アプリケーション (クライアント) ID`と`ディレクトリ (テナント) ID`をコピーします。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/33a01f09-1934-4a7c-bbe3-63b20324d0ba.png)

`証明書とシークレット` -> `新しいクライアント シークレット`を選択します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/f9f4aa73-69dd-4879-9bbd-11e19edcb066.png)

`追加`を選択します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/81903f00-d518-458c-87ee-27b3d164bb76.png)

表示されているシークレットの値をコピーします。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/c492305b-8dd5-469e-9c6f-2f03554a7f4e.png)


### 2. Azure Event Hubs にてアプリに対して`Azure Event Hubs のデータ受信者`権限を付与

Azure Event Hubs 名前空間（あるいは、 Azure Event Hubs インスタンス）にて、`アクセス制御 (IAM)` -> `+ 追加` -> `ロールの割り当ての追加`を選択します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/3d7d86c6-a9b1-45a7-9151-0f0ca6bb8468.png)

`Azure Event Hubs のデータ受信者`権限を作成したアプリに付与します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/b3893c24-9029-4f27-91ee-65367945cecf.png)


### 3. Databriks にて Azure Event Hubs からデータを取得

Azure Event Hubs に接続するための設定値をセットします。

```python
# Azure Event Hubs に接続するための設定値をセット
eh_namespace_name = "kafka-test-01"  ## Azure Event Hubs の名前空間をセット
ehi_name = topic_name = "dbx-to-eh-01"  ## Azure Event Hubs のEvent Hubs のインスタンスをセット

tenant_id = "095ca098-XXXX-XXX-XXX-XXX"  # Azure Entra ID の Tenant ID をセット
client_id = "214cee56-7fd6-418d-bd21-aded919113b4e"  # Azure Entra ID の Application (クライアント) ID をセット
client_secret = "pIj8Q~fPMueU3YHXzyYmXXXXXX"  # Azure Entra ID の Application (クライアント) Secret をセット
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/79dd7520-4a5c-4a3d-bd47-b3f3b07bdc75.png)

Azure Event Hubs からデータを取得できることを確認します。

```python
# Databricks から Azure Event Hubs に設定する際の設定をセット
bootstrap_servers = f"{eh_namespace_name}.servicebus.windows.net"
sasl_config = ""
sasl_config = (
    "kafkashaded.org.apache.kafka.common.security.oauthbearer.OAuthBearerLoginModule required "
    f'clientId="{client_id}" '
    f'clientSecret="{client_secret}" '
    f'scope="https://{bootstrap_servers}/.default" ssl.protocol="SSL";'
)

kafka_options = {
    "kafka.bootstrap.servers": f"{bootstrap_servers}:9093",
    "kafka.sasl.jaas.config": sasl_config,
    "kafka.sasl.oauthbearer.token.endpoint.url": f"https://login.microsoft.com/{tenant_id}/oauth2/v2.0/token",
    "kafka.security.protocol": "SASL_SSL",
    "kafka.sasl.mechanism": "OAUTHBEARER",
    "kafka.sasl.login.callback.handler.class": "kafkashaded.org.apache.kafka.common.security.oauthbearer.secured.OAuthBearerLoginCallbackHandler",
}

# Azure Event Hubs からデータを取得
from pyspark.sql.functions import col, expr, from_json

df = (
    spark.readStream.format("kafka")
    .options(**kafka_options)
    .option("subscribe", topic_name)
    .option("startingOffsets", "earliest")  # 検証目的であるため earliest に設定
    .load()
)
df = df.withColumn(
    "decoded_value",
    col("value").cast("string"),
)

# データを表示
display(df)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/24031/d2873082-4a60-4b9e-b74c-ea24d7c9cc62.png)