🎐

AWS Lambda カスタムイメージを使って PostgreSQL に接続する

2023/12/05に公開

はじめに

AWS Lambda から Python を使って PostgreSQL に接続するときに、 psycopg2 を使うことが多いと思います。
その場合、ZIP にパッケージングして Lambda にアップロードする必要があります。
今までは、 Serverless Framework を使ったり、 lambci/docker-lambda を使ってローカルでパッケージングしていました。

普段、 Terraform を使って Lambda などの環境を構築する際に、 ZIP ファイルを予めリポジトリに入れられれば管理が簡単ですが、ファイルが大きい場合デプロイ時に ZIP にパッケージングして S3 へアップロードするのが面倒です。
そこで、 AWS Lambda カスタムイメージを使って PostgreSQL に接続する方法を試してみました。

準備

AWS Lambda イメージ

AWS Lambda のベースイメージが ECR パブリックギャラリーに公開されています。
このイメージをベースに、 psycopg2 をインストールしたイメージを作成します。

イメージの作成

Dockerfile

次のような Dockerfile を作成します。

この Dockerfile では、特定のバージョンの AWS Lambda の Pythonイメージ( public.ecr.aws/lambda/python:3.11.2023.11.01.21-x86_64 )を使用しています。
x86_64 アーキテクチャで Lambda 関数を実行したかったので特定のイメージを選択しました。

FROM public.ecr.aws/lambda/python:3.11.2023.11.01.21-x86_64
ENV PGVERSION 13.13

RUN yum update -y \
  && yum install -y tar gzip gcc python3-devel python3-pip python3-setuptools libxml2-devel libxslt-devel readline-devel uuid-devel openssl \
  && yum clean all

RUN curl -O https://ftp.postgresql.org/pub/source/v${PGVERSION}/postgresql-${PGVERSION}.tar.gz \
  && tar -xvf postgresql-${PGVERSION}.tar.gz \
  && cd postgresql-${PGVERSION} \
  && ./configure \
  && make -j4 \
  && make install \
  && cd .. \
  && rm -rf postgresql-${PGVERSION} postgresql-${PGVERSION}.tar.gz

COPY src/*.py $LAMBDA_TASK_ROOT
COPY src/requirements.txt $LAMBDA_TASK_ROOT

RUN pip install --upgrade pip \
  && pip install -r requirements.txt

CMD ["handler.main"]

AWS Lambda の標準イメージには PostgreSQL 13 のパッケージが含まれていないため、 ソースコード から直接ビルドし、インストールする方法を採用しました。
この方法は特定のバージョンの要件を満たす必要がある場合に利用することができます。

requirements.txt には、 psycopg2 を記載しています。

psycopg2-binary

イメージのビルドと更新

ビルドしたイメージを ECR にプッシュするための Makefile を作成します。

latest タグでプッシュしたら、 Lambda 関数の更新されると思ったのですが自動では更新されません。aws lambda update-function-code を実行して明示的に image-uri を指定して更新する必要があります。

ECR にプッシュしたときに Lambda 関数を更新するために update-function-code コマンドを実行するようにしています。

Makefile

APP_NAME = lambda/lambda-postgresql
FUNCTION_NAME = lambda-postgresql-sandbox
APP_VERSION ?=

AWS_ECR_ACCOUNT_ID ?=
AWS_ECR_REGION ?=
AWS_ECR_REPO = $(APP_NAME)

TAG ?= latest


.PHONY : docker/build docker/push docker/run/bash

docker/build :
	docker build -t $(APP_NAME):$(APP_VERSION) .


docker/push : docker/build
	aws ecr get-login-password --region $(AWS_ECR_REGION) | docker login --username AWS --password-stdin $(AWS_ECR_ACCOUNT_ID).dkr.ecr.$(AWS_ECR_REGION).amazonaws.com
	docker tag $(APP_NAME):$(APP_VERSION) $(AWS_ECR_ACCOUNT_ID).dkr.ecr.$(AWS_ECR_REGION).amazonaws.com/$(AWS_ECR_REPO):$(TAG)
	docker push $(AWS_ECR_ACCOUNT_ID).dkr.ecr.$(AWS_ECR_REGION).amazonaws.com/$(AWS_ECR_REPO):$(TAG)
	aws lambda update-function-code --function-name $(FUNCTION_NAME) --image-uri $(AWS_ECR_ACCOUNT_ID).dkr.ecr.$(AWS_ECR_REGION).amazonaws.com/$(AWS_ECR_REPO):$(TAG) --publish || true

Makefile を使用することで、複数のステップを含むビルドプロセスを単一のコマンドで実行できます。この方法は、ビルドからデプロイまでを標準化し CD を実現する際にも有用です。

実行例

以下のように実行すると、イメージがビルドされ ECR にプッシュされます。
同時に Lambda 関数が更新されるようになりますので、かなり更新が楽になります。

$ export AWS_ACCESS_KEY_ID="ASIxxx"
export AWS_SECRET_ACCESS_KEY="Aooxxx"
export AWS_SESSION_TOKEN="IQoxxx"

$ make docker/push AWS_ECR_ACCOUNT_ID=123456789012 AWS_ECR_REGION=ap-northeast-1 APP_VERSION=latest

Lambda 関数の中身と実行

コンテナ内に実行する以下の Python のコードを配置しています。
中身は、 psycopg2 を使って PostgreSQL に接続して、 information_schema.columns からテーブル名とカラム名を取得してログに出力します。

import logging
import os

import psycopg2

logger = logging.getLogger()
logger.setLevel(logging.INFO)


def execute_sql(database_url, sql):
  rows = None
  with psycopg2.connect(database_url) as conn:
    with conn.cursor() as cursor:
      cursor.execute(sql)
      if 'SELECT' in cursor.statusmessage:
        rows = cursor.fetchall()
  return rows


def column_all_sql():
  return """
    SELECT
      c.table_name,
      c.column_name
    FROM
      information_schema.columns AS c
    WHERE
      c.table_schema = 'public'
    ORDER BY
      c.table_name,
      c.column_name;
  """


def main(event, context):
  logger.info(event)

  database_url = os.environ.get('DATABASE_URL')

  if database_url is None:
    raise Exception('Not found DATABASE_URL environment value')

  rows = execute_sql(database_url, column_all_sql())
  logger.info(rows)
  logger.info(rows[0])

Lambda 関数と PostgreSQL の間のセキュアな接続を確保するためには、 Secret Manager を通じてデータベース接続情報を渡すほうがよいです。
今回、簡単のために環境変数に直接接続情報を設定しています。

実行ログ

今回 RDS に PostgreSQL 13 を起動して dvdrental データベースを作成しています。
DATABASE_URL には、 その作成した RDS のエンドポイントを指定しています。セキュリティグループの設定などは、適宜行ってください。

lambda 実行してログを確認すると、以下のようにテーブル名とカラム名が取得できていることが確認できます。

[INFO]	2023-11-25T16:46:38.195Z	fa465df2-479f-45f7-9827-d60d37d11857	[('actor', 'actor_id'), ('actor', 'first_name'), ('actor', 'last_name'), ('actor', 'last_update') ...
[INFO] 2023-11-25T16:46:38.195Z fa465df2-479f-45f7-9827-d60d37d11857 ('actor', 'actor_id')

実際に作成したものは以下にあります。

まとめ

AWS Lambda カスタムイメージを使って PostgreSQL に接続する方法を試してみました。
Terraform の環境構築と Lambda 関数のデプロイを分けることができるので、便利だと思いました。

Linc'well, inc.

Discussion