【AWS Glue】AWS GlueのBigQueryコネクタを使って直接BigQueryへデータを書き込んでみよう

import sys
from awsglue.transforms import *
from awsglue.utils import getResolvedOptions
from pyspark.context import SparkContext
from awsglue.context import GlueContext
from awsglue.job import Job
from awsglue.dynamicframe import DynamicFrame
from pyspark.conf import SparkConf

# ジョブパラメータ
args = getResolvedOptions(sys.argv, [
    'JOB_NAME',
    'bigquery_connection_name',
    'source_database',
    'source_table',
    'target_dataset',
    'target_table',
    'gcp_project_id'
])

# Spark設定
conf = SparkConf()
conf.set('spark.sql.catalog.s3tablesbucket', 'org.apache.iceberg.spark.SparkCatalog')
conf.set('spark.sql.catalog.s3tablesbucket.catalog-impl', 'software.amazon.s3tables.iceberg.S3TablesCatalog')
conf.set('spark.sql.catalog.s3tablesbucket.warehouse', 'arn:aws:s3tables:<region>:<account-id>:bucket/<bucket-name>')
conf.set('spark.sql.extensions', 'org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions')

# SparkとGlueコンテキストの初期化
sc = SparkContext(conf=conf)
glueContext = GlueContext(sc)
spark = glueContext.spark_session
job = Job(glueContext)
job.init(args['JOB_NAME'], args)

print("=== Starting S3 Tables to BigQuery ETL Job ===")

try:
    # S3TablesCatalogを使用してS3 Tablesからデータを読み取り
    source_table_name = f"s3tablesbucket.{args['source_database']}.{args['source_table']}"
    print(f"Reading from S3 Tables: {source_table_name}")

    # SQLを使用してデータを読み取り
    df_source = spark.sql(f"SELECT * FROM {source_table_name}")

    # スキーマとレコード数を表示
    print("Source data schema:")
    df_source.printSchema()

    record_count = df_source.count()
    print(f"Number of records to transfer: {record_count}")

    # Spark DataFrameをGlue用のDynamicFrameに変換
    dynamic_frame = DynamicFrame.fromDF(df_source, glueContext, "dynamic_frame")

    # BigQueryに書き込み
    print("Writing data to BigQuery...")

    bigquery_write = glueContext.write_dynamic_frame.from_options(
        frame=dynamic_frame,
        connection_type="bigquery",
        connection_options={
            "connectionName": args['bigquery_connection_name'],
            "parentProject": args['gcp_project_id'],
            "table": f"{args['target_dataset']}.{args['target_table']}",
            "writeMethod": "direct"
        },
        transformation_ctx="bigquery_write"
    )

    print(f"Successfully transferred {record_count} records to BigQuery")

except Exception as e:
    print(f"Error occurred during ETL process: {str(e)}")
    raise e
finally:
    job.commit()

# スクリプトをS3にアップロード
aws s3 cp glue_s3tables_to_bigquery.py s3://<your-bucket>/scripts/

# Glueジョブを作成
aws glue create-job --cli-input-json file://glue_job_definition.json

"=== Starting S3 Tables to BigQuery ETL Job ===
Source: cm_kawanago_tablebucket_namespace.daily_sales
Target: <your-gcp-project-id>.cm_kawanago.daily_sales_from_aws
Reading from S3 Tables: s3tablesbucket.cm_kawanago_tablebucket_namespace.daily_sales"

"Source data schema:"
"root
 |-- sale_date: date (nullable = true)
 |-- product_category: string (nullable = true)
 |-- sales_amount: double (nullable = true)"

"Number of records to transfer: 9
Sample data (first 10 rows):"

"
+----------+----------------+------------+
|sale_date |product_category|sales_amount|
+----------+----------------+------------+
|2024-01-15|Laptop          |900.0       |
|2024-01-15|Monitor         |250.0       |
|2024-01-16|Laptop          |1350.0      |
|2024-02-01|Monitor         |300.0       |
|2024-02-01|Keyboard        |60.0        |
|2024-02-02|Mouse           |25.0        |
|2024-02-02|Laptop          |1050.0      |
|2024-02-03|Laptop          |1200.0      |
|2024-02-03|Monitor         |375.0       |
+----------+----------------+------------+
"

"Writing data to BigQuery..."
"Successfully transferred 9 records to BigQuery
=== ETL Job Completed Successfully ==="

"Running autoDebugger shutdown hook."