Databrick REST APIによりWorkspaceオブジェクト（ノートブック・Experiments・フォルダ）の存在チェックを行う方法 #Python

概要

Databrick REST APIによりWorkspaceオブジェクト（ノートブック・Experiments・フォルダ）の存在チェックを行う方法を共有します。dbutilsの機能にはないため、REST APIで実施する必要があります。

詳細は下記のGithub pagesのページをご確認ください。

コードを実行したい方は、下記のdbcファイルを取り込んでください。

https://github.com/manabian-/databricks_tecks_for_qiita/blob/main/tecks/check_ws_obj_by_rest_api/dbc/check_ws_obj_by_rest_api.dbc

実行環境

databricks runtime: 8.3.x-cpu-ml-scala2.12
Python version: 3.8.8
pyspark version: 3.1.2.dev0

手順

1. シークレットにDatabricksトークンとDatabricks Workspace URLを格納して、変数にセット

シークレットにDatabricksトークンとDatabricks Workspace URLを格納する方法は下記の記事を参考に実施してください。

Databrick REST APIによりDatabricksトークンを作成後にシークレットとして登録する方法 - Qiita

scope_name = "databricks_mlops" # scope名をセット
databricks_token = "databricks_token" # DatabricksトークンにおけるsecreatのKey名をセット
databricks_url = "databricks_url" # DatabricksのURLをsecretsのKey名をセット

# シークレットからトークンとDatabrikcs WorkspaceのURLをセット
token  = dbutils.secrets.get(scope_name, databricks_token)
db_url = dbutils.secrets.get(scope_name, databricks_url)

2. 関数を定義

import requests

def check_ws_obj_exist(path, token, db_url,raise_error='False'):
    """Databricks Workspaceにおけるオブジェクトの存在チェック
        Databricks Workspaceにおけるオブジェクトの存在チェックを、Databricks REST APIにより行う関数です。

        Args:
            path (string): ノートブック等のWorkspaceオブジェクトのパスを指定
            token (string): REST API実行する際に利用するDatabricksトークンを指定
            db_url (string): REST API実行対象のDatabricks Workspace URLを指定
            raise_error (boolean): 正常終了とする場合にはTrue、異常終了にする場合にはFalse

        Returns:
            None

        Raises:
            Exception: Notebook does not exist (raise_errorがTrueであり、オブジェクトが存在しない場合)

        Yields:
            None

        Examples:
            >>> chek_notebook_existence(path, False)
               Notebook does not exist

            >>> chek_notebook_existence(path, True)
               Exception: Notebook does not exist

        Note:
            Nothing

        """    
    response = requests.get(
        f'{db_url}/api/2.0/workspace/list',
        headers={'Authorization': f'Bearer {token}'},
        json={
            'path': path,
        },
    )

    if response.status_code == 200:
        print(f'Workspace objects exist')
    elif response.status_code == 404:
        if response.json().get('error_code') == 'RESOURCE_DOES_NOT_EXIST':
            if raise_error == True:
                raise Exception(f'Workspace objects does not exist')
            else:
                print(f'Workspace objects does not exist')
        else:
            print(response.text)
    elif response.status_code == 403:
        print(response.text)
    else:
        print('Error geting the job: {0}: {1}'.format(response.json()['error_code'],response.json()['message']))

3. 定義した関数によりDatabricks Workspaceオブジェクトの存在チェック（存在しなくても正常終了）

# 存在チェックしたいノートブックのパスを指定
path = "/qiita/aaa"

check_ws_obj_exist(path, token, db_url, False)

4. 定義した関数によりDatabricks Workspaceオブジェクトの存在チェック（存在しない場合に異常終了）

# 存在チェックしたいノートブックのパスを指定
path = "/qiita/aaa"

check_ws_obj_exist(path, token, db_url, True)