Databricks SDK for Pythonを使ったシステムスキーマの有効化

# はじめに
Databricksのシステムスキーマの有効化について、以下のドキュメントでは[SystemSchemas API](https://docs.databricks.com/api/azure/workspace/systemschemas)を使った方法が記載されています。

https://learn.microsoft.com/ja-jp/azure/databricks/admin/system-tables/

上記の方法はもちろん問題なく動作するのですが、APIを呼び出すためのパーソナルアクセストークン（PAT）の発行などの手間が生じます。

本記事では、**Databricksノートブック上でDatabricks SDK for Python**を用いてシステムスキーマを有効化する方法を紹介します。この方法では、PATを発行することなくシステムスキーマを有効化できます。

## 前提条件、動作確認環境など
- 処理を実行するユーザーはアカウント管理者の権限を持っている必要があります
- Databricks Runtime 14.3 LTS、シングルノードで動作確認を行っています
    - 他のバージョンでもおそらく問題なく動作すると思います

## 参考リンク
https://databricks-sdk-py.readthedocs.io/en/latest/workspace/catalog/system_schemas.html

https://databricks-sdk-py.readthedocs.io/en/latest/workspace/catalog/metastores.html

# システムスキーマを1つずつ有効化
まずは分かりやすさのためにステップバイステップで1つずつシステムスキーマを有効化する方法を記載します。
（次のセクションで、まとめてシステムスキーマを有効化する方法を記載していますので、最初からそちらを使って頂いてももちろん問題ありません）

## 1. Databricks SDK for Pythonの最新化
Databricks SDK for Pythonを最新化し、Pythonを再起動します。

```python
%pip install databricks-sdk --upgrade -q
dbutils.library.restartPython()
```

## 2. 更新前のシステムスキーマをリストアップ
ワークスペースが関連付けられているメタストアを取得し、メタストアIDを引数としてシステムスキーマをリストアップします。

```python
from databricks.sdk import WorkspaceClient

w = WorkspaceClient()
current_metastore = w.metastores.current()
w.system_schemas.list(current_metastore.metastore_id)
```

以下のような結果が表示されます。`state`が`SystemSchemaInfoState.AVAILABLE`になっているものが有効化可能なシステムスキーマです。

```python:更新前のシステムスキーマ一覧
[SystemSchemaInfo(schema='access', state=<SystemSchemaInfoState.ENABLE_COMPLETED: 'ENABLE_COMPLETED'>),
 SystemSchemaInfo(schema='billing', state=<SystemSchemaInfoState.ENABLE_COMPLETED: 'ENABLE_COMPLETED'>),
 SystemSchemaInfo(schema='compute', state=<SystemSchemaInfoState.ENABLE_COMPLETED: 'ENABLE_COMPLETED'>),
 SystemSchemaInfo(schema='marketplace', state=<SystemSchemaInfoState.ENABLE_COMPLETED: 'ENABLE_COMPLETED'>),
 SystemSchemaInfo(schema='workflow', state=<SystemSchemaInfoState.AVAILABLE: 'AVAILABLE'>),
 SystemSchemaInfo(schema='operational_data', state=<SystemSchemaInfoState.UNAVAILABLE: 'UNAVAILABLE'>),
 SystemSchemaInfo(schema='lineage', state=<SystemSchemaInfoState.UNAVAILABLE: 'UNAVAILABLE'>),
 SystemSchemaInfo(schema='hms_to_uc_migration', state=<SystemSchemaInfoState.UNAVAILABLE: 'UNAVAILABLE'>),
 SystemSchemaInfo(schema='information_schema', state=<SystemSchemaInfoState.ENABLE_COMPLETED: 'ENABLE_COMPLETED'>)]
```

## 3. システムスキーマの有効化
上記の例では`workflow`が有効化可能なので、こちらを直接指定して有効化してみます。

```python
w.system_schemas.enable(current_metastore_id, "workflow")
```

カタログエクスプローラーで`system`カタログを見ると、有効化されていることが確認できます。複数のシステムスキーマを有効化したい場合は、上記処理をシステムスキーマ名を変えて複数回実行します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/284925/2ba309ca-e8b9-a5d1-4851-41506cd6e41c.png)

# システムスキーマをまとめて有効化
`state`が`SystemSchemaInfoState.AVAILABLE`なシステムスキーマをまとめて有効化したい場合には、以下のコードを実行します。

```python
%pip install databricks-sdk --upgrade -q
dbutils.library.restartPython()
```

```python
from typing import List
from databricks.sdk import WorkspaceClient
from databricks.sdk.service.catalog import SystemSchemaInfo, SystemSchemaInfoState

def get_all_schemas(workspace_client: WorkspaceClient, metastore_id: str) -> List[SystemSchemaInfo]:
    """
    指定されたメタストアIDのすべてのシステムスキーマを取得します。

    :param workspace_client: DatabricksのWorkspaceClientインスタンス
    :param metastore_id: メタストアのID
    :return: SystemSchemaInfoのリスト
    """
    return workspace_client.system_schemas.list(metastore_id)

def filter_available_schemas(schemas: List[SystemSchemaInfo]) -> List[SystemSchemaInfo]:
    """
    利用可能な状態のシステムスキーマのみをフィルタリングします。

    :param schemas: SystemSchemaInfoのリスト
    :return: 利用可能な状態のSystemSchemaInfoのリスト
    """
    return [schema_info for schema_info in schemas if schema_info.state == SystemSchemaInfoState.AVAILABLE]

def enable_schemas(workspace_client: WorkspaceClient, metastore_id: str, schemas: List[SystemSchemaInfo]) -> None:
    """
    指定されたリストのシステムスキーマをすべて有効化します。
    エラーが発生した場合はキャッチして情報を出力します。

    :param workspace_client: DatabricksのWorkspaceClientインスタンス
    :param metastore_id: メタストアのID
    :param schemas: 有効化するSystemSchemaInfoのリスト
    """
    for schema_info in schemas:
        try:
            workspace_client.system_schemas.enable(metastore_id, schema_info.schema)
            print(f"システムスキーマ '{schema_info.schema}' を有効化しました。")
        except Exception as e:
            print(f"システムスキーマ '{schema_info.schema}' の有効化中にエラーが発生しました: {str(e)}")

def main() -> None:
    """
    メイン関数。以下の手順を実行します:
    1. クライアントの初期化、現在のメタストアIDを取得
    2. システムスキーマのリストを取得
    3. 有効化可能なシステムスキーマに絞り込み
    4. 利用可能なすべてのスキーマを有効化
    5. 有効化後のシステムスキーマのリストを再取得
    """
    # クライアントの初期化、現在のメタストアIDを取得
    w = WorkspaceClient()
    current_metastore = w.metastores.current()
    current_metastore_id = current_metastore.metastore_id

    # システムスキーマのリストを取得
    all_schemas = get_all_schemas(w, current_metastore_id)
    print(f"All schemas before update: {all_schemas}")

    # 有効化可能なシステムスキーマに絞り込み
    available_schemas = filter_available_schemas(all_schemas)
    print(f"Available schemas: {available_schemas}")

    # 利用可能なすべてのスキーマを有効化
    enable_schemas(w, current_metastore_id, available_schemas)

    # 有効化後のシステムスキーマのリストを再取得
    all_schemas_after_update = get_all_schemas(w, current_metastore_id)
    print(f"All schemas after update: {all_schemas_after_update}")
```

```python
# 上記を実行
main()
```

以上です。