S3メタデータからAthena UDFを使用して署名付きURLを生成してみた

# S3バケットの作成
aws s3 mb s3://cm-watanabe-metadata \
    --region us-east-1 

# S3メタデータテーブル用のS3テーブルバケットの作成
aws s3tables create-table-bucket \
    --name cm-watanabe-table-bucket \
    --region us-east-1 \

# S3メタデータ統合の設定
aws s3api create-bucket-metadata-table-configuration \
    --bucket cm-watanabe-metadata \
    --region us-east-1 \
    --metadata-table-configuration '{
        "S3TablesDestination": {
            "TableBucketArn": "arn:aws:s3tables:us-east-1:<account-id>:bucket/cm-watanabe-table-bucket",
            "TableName": "cm_watanabe_meta_table"
        }
    }'    

# ファイル配置
aws s3 cp ./clanyan1.jpg s3://cm-watanabe-metadata \
    --metadata '{"location":"hibiya"}' \

aws s3 cp ./clanyan2.jpg s3://cm-watanabe-metadata \
    --metadata '{"location":"osaka"}' \

aws s3 cp ./clanyan3.jpg s3://cm-watanabe-metadata \
    --metadata '{"location":"fukuoka"}' \

.
├── README.md
├── athena_udf_lambda
│   ├── __init__.py
│   ├── app.py
│   └── requirements.txt
├── samconfig.toml
├── template.yaml

AWSTemplateFormatVersion: '2010-09-09'
Transform: AWS::Serverless-2016-10-31
Description: >
  athena-udf

Globals:
  Function:
    Timeout: 300
Parameters:
  PythonVersion:
    Type: String
    Description: Python version
    AllowedValues:
      - "3.9"
      - "3.10"
      - "3.11"
      - "3.12"
    ConstraintDescription: Must be one of the allowed values.

Resources:
  AthenaUDFFunction:
    Type: AWS::Serverless::Function 
    Properties:
      CodeUri: athena_udf_lambda/
      FunctionName: athena_s3_presigned_url_udf
      Handler: app.lambda_handler
      Runtime: !Sub "python${PythonVersion}"
      MemorySize: 256
      Architectures:
        - x86_64
      Policies:
        - AmazonS3ReadOnlyAccess
Outputs:
  AthenaUDFFunction:
    Description: "Athena UDF Function ARN"
    Value: !GetAtt AthenaUDFFunction.Arn

# More information about the configuration file can be found here:
# https://docs.aws.amazon.com/serverless-application-model/latest/developerguide/serverless-sam-cli-config.html
version = 0.1

[default.global.parameters]
stack_name = "athena-udf"

[default.build.parameters]
cached = true
parallel = true

[default.validate.parameters]
lint = true

[default.deploy.parameters]
capabilities = "CAPABILITY_IAM"
confirm_changeset = true
resolve_s3 = true
s3_prefix = "athena-udf"
region = "us-east-1"
image_repositories = []

[default.package.parameters]
resolve_s3 = true

[default.sync.parameters]
watch = true

[default.local_start_api.parameters]
warm_containers = "EAGER"

[default.local_start_lambda.parameters]
warm_containers = "EAGER"

import boto3
import athena_udf
from botocore.exceptions import ClientError

# S3クライアントを初期化してAmazon S3とやり取りします
s3_client = boto3.client('s3')

class PresignedUrlUDF(athena_udf.BaseAthenaUDF):
    """
    PresignedUrlUDFクラスはBaseAthenaUDFから継承され、AWS Athenaのクエリを処理し、
    S3オブジェクトに対する事前署名付きURLを返します。

    主なメソッド:
        - handle_athena_record: Athenaクエリの入力を受け取り、事前署名付きURLを生成して返します。
    """

    @staticmethod
    def handle_athena_record(input_schema, output_schema, arguments):
        """
        Athena UDFのレコードを処理し、S3オブジェクトの事前署名付きURLを生成します。

        引数:
            input_schema (list): UDFの入力スキーマ（このロジックでは使用されません）。
            output_schema (list): UDFの出力スキーマ（このロジックでは使用されません）。
            arguments (list): UDFに渡される引数のリスト。引数は以下の通り:
                - arguments[0]: bucket_name (str) - S3バケットの名前。
                - arguments[1]: object_key (str) - バケット内のオブジェクトのキー。
                - arguments[2] (オプション): expiration (int) - URLの有効期限（秒単位）。デフォルトは3600秒（1時間）。

        戻り値:
            str: S3オブジェクトの事前署名付きURL。

        例外:
            ValueError: bucket_name または object_key が不足している場合に発生します。
        """

        # argumentsからbucket_nameとobject_keyを取得
        bucket_name = arguments[0]
        object_key = arguments[1]
        expiration = int(arguments[2]) if len(
            arguments) > 2 else 3600  # デフォルトは1時間

        # bucket_nameとobject_keyが空でないことを確認
        if not bucket_name or not object_key:
            raise ValueError(
                "Both bucket_name and object_key must be provided.")

        # presigned URLを生成
        presigned_url = generate_presigned_url(
            bucket_name, object_key, expiration)

        return presigned_url

def generate_presigned_url(bucket_name, object_key, expiration):
    """
    S3オブジェクトのダウンロード用に事前署名付きURLを生成します。

    引数:
        bucket_name (str): S3バケットの名前。
        object_key (str): バケット内のオブジェクトのキー。
        expiration (int): URLの有効期限（秒単位）。

    戻り値:
        str: S3オブジェクトの事前署名付きURL。

    例外:
        ValueError: 事前署名付きURLの生成中にエラーが発生した場合。
    """
    try:
        # S3からオブジェクトをダウンロードするためのpresigned URLを生成
        url = s3_client.generate_presigned_url(
            'get_object',
            Params={'Bucket': bucket_name, 'Key': object_key},
            ExpiresIn=expiration
        )
        return url
    except ClientError as e:
        raise ValueError(f"Error generating presigned URL: {str(e)}")

# UDFの`lambda_handler`関数にlambda_handlerを割り当て
lambda_handler = PresignedUrlUDF().lambda_handler

athena_udf

sam build --parameter-overrides PythonVersion=3.11
sam deploy --parameter-overrides PythonVersion=3.11

SELECT * FROM "cm_watanabe_meta_table" WHERE user_metadata['location']='hibiya'

USING EXTERNAL FUNCTION s3_presigned_url(bucket VARCHAR, key VARCHAR)
RETURNS VARCHAR
LAMBDA 'arn:aws:lambda:us-east-1:<account-id>:function:athena_s3_presigned_url_udf'
SELECT 
    bucket,
    key,
    s3_presigned_url(bucket, key) as presigned_url
FROM "aws_s3_metadata"."cm_watanabe_meta_table"
WHERE user_metadata['location'] = 'hibiya'