Treasure Data の Custom Scripts で任意のライブラリを読み込む

Custom Scripts で未インストールのライブラリを使いたい場合のメモです。

例えば、以下のようにファイルが配置されている時、

├ py_scripts
│ └ sample.py
└ sample.dig

ワークフローの設定ファイルには以下のように記述し、

timezone: Asia/Tokyo

+task1:
  py>: py_scripts.sample.sample_sleep
  docker:
    image: "digdag/digdag-python:3.7"

sample.py の中で以下のようにインストールします。

import sys
import subprocess

subprocess.call([sys.executable, '-m', 'pip', 'install', 'tqdm'])

def sample_sleep():
  from tqdm import tqdm
  import time

  for i in tqdm(range(100)):
    time.sleep(0.1)

参考になった記事

Introduction to Custom Scripts
Adding a Custom Python Script to your Workflow