DatabricksでSQLウェアハウスを構築して、SQLAlchemy+AlembicでDB Migrationを行う

はじめに

DatabricksのSQLウェアハウスは、Databricks Appsのデータベースとして使用できるものです。そのため、アプリケーションと統合されたデータベースマイグレーションを行い、データベースのバージョン管理を行えることを目指します。

本記事では、以下の内容について解説します：

• DatabricksでSQLウェアハウスを構築する方法
• SQLAlchemyとAlembicを使用してデータベースマイグレーションを行う方法

前段

• Azure DatabricksでDatabricks Appsを検証するための環境構築
• Databricks Apps でカスタムアプリのデプロイ
• FastAPIにSPAアプリケーションを統合する
• Databricks AppsにSPAアプリを含めたカスタムアプリを作成する
• Databricks WorkspaceとGitHubを連携してDatabricks Appsにデプロイする
• Databricks WebTerminalを有効化する

手順

DatabricksでSQL Workspaceを構築する

1. 画面左メニューからSQL>SQLウェアハウスを選択、画面上部の「SQLウェアハウスを作成」ボタンを押下
   

2. SQLウェアハウスの名前を入力し、他はデフォルト設定のまま、「作成」ボタンを押下
   

3. 権限管理設定はデフォルト設定のまま、×ボタンを押下
   

4. SQLウェアハウスの一覧画面から作成したSQLウェアハウスを選択し、概要を確認する
   

SQLAlchemy+AlembicでDB Migrationを行う(アプリ側設定は概要のみ記載)

※ コードベースは過去の記事のものを使用しています
※ 作業環境は過去の記事で紹介したWeb Terminalを使用しています

1. SQLAlchemyとAlembicをインストールする

   pip install sqlalchemy alembic databricks-sql-connector
   

2. models.pyを作成し、テーブルのモデルを作成する

   from sqlalchemy import Column, Integer, String, DateTime
   from sqlalchemy.ext.declarative import declarative_base
   from datetime import datetime
   
   Base = declarative_base()
   
   class Item(Base):
   __tablename__ = "items"
   
   id = Column(Integer, primary_key=True, index=True)
   item_id = Column(String, index=True)
   query = Column(String)
   created_at = Column(DateTime, default=datetime.utcnow)
   updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
   

3. alembicの設定ファイルを作成する

   alembic init alembic
   

4. alembic.iniを編集する

   # 以下の環境変数が必要です：
   # DATABRICKS_SERVER_HOSTNAME - Databricksのホスト名
   # DATABRICKS_HTTP_PATH - SQLウェアハウスのHTTPパス
   # DATABRICKS_TOKEN - アクセストークン
   sqlalchemy.url = databricks://token:%(DATABRICKS_TOKEN_PROD)s@%(DATABRICKS_SERVER_HOSTNAME_PROD)s?http_path=%(DATABRICKS_HTTP_PATH_PROD)s&catalog=%(DATABRICKS_CATALOG_PROD)s&schema=%(DATABRICKS_SCHEMA_PROD)s
   

5. env.pyを編集する
   engine_from_configだとエラーになったので、create_engineに書き換えました。

   connectable = create_engine(DATABASE_URL, connect_args={"verify": False})
   

6. SQLウェアハウスの一覧画面から作成したSQLウェアハウスを選択し、接続の詳細からサーバのホスト名、HTTPパスを取得する
   

7. 同画面からPythonを選択し、画面内「新規トークンを作成」ボタンを押下し、トークンを作成、コピーする
   

8. 取得した内容を環境変数に設定する

   export ENV=production
   export DATABRICKS_SERVER_HOSTNAME_PROD=<サーバのホスト名>
   export DATABRICKS_HTTP_PATH_PROD=<HTTPパス>
   export DATABRICKS_TOKEN_PROD=<トークン>
   export DATABRICKS_CATALOG_PROD=<カタログ名>(初期値はワークスペース名)
   export DATABRICKS_SCHEMA_PROD=<スキーマ名>(初期値はdefault)
   

9. コマンドを実行して、マイグレーションを実行する

   alembic revision --autogenerate -m "Initial migration"
   alembic upgrade head
   

   実行すると、なぜかエラーになりますが、マイグレーションは正常に実行されています。

   sqlalchemy.exc.DatabaseError: (databricks.sql.exc.ServerOperationError) [UC_COMMAND_NOT_SUPPORTED.WITHOUT_RECOMMENDATION] The command(s): CreateIndex are not supported in Unity Catalog.
   

10. マイグレーションの確認
    画面左メニューのカタログから、指定したカタログ、スキーマを選択すると、以下2テーブルが作成されていることが確認できます

    • alembic_version
    • items

    

まとめ

本記事では、以下の内容について解説しました：

1. DatabricksでのSQLウェアハウスの構築

   • SQLウェアハウスの作成と基本設定

2. SQLAlchemy + Alembicを使用したDBマイグレーション

   • 必要なパッケージのインストール（sqlalchemy, alembic, databricks-sql-connector）
   • モデル定義とマイグレーション設定
   • 環境変数を使用した接続設定
   • マイグレーションの実行手順

3. 注意点

   • Unity Catalogの制限により、CreateIndexコマンドはサポートされていないが、マイグレーション自体は正常に実行される
   • 接続設定には環境変数（ホスト名、HTTPパス、トークン、カタログ名、スキーマ名）が必要

参考

• Databricks SQLAlchemy