【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()

{
    "Name": "s3tables-to-bigquery-daily-sales",
    "Description": "S3 Tables から BigQuery へのデータ連携ジョブ",
    "Role": "arn:aws:iam::<account-id>:role/<your-glue-role>",
    "ExecutionProperty": {
        "MaxConcurrentRuns": 1
    },
    "Command": {
        "Name": "glueetl",
        "ScriptLocation": "s3://<your-bucket>/scripts/glue_s3tables_to_bigquery.py",
        "PythonVersion": "3"
    },
    "DefaultArguments": {
        "--enable-metrics": "true",
        "--enable-spark-ui": "true",
        "--spark-event-logs-path": "s3://<your-bucket>/spark-logs/",
        "--enable-job-insights": "true",
        "--enable-continuous-cloudwatch-log": "true",
        "--job-bookmark-option": "job-bookmark-disable",
        "--job-language": "python",
        "--datalake-formats": "iceberg",
        "--additional-python-modules": "pyarrow==14.0.1",
        "--extra-jars": "s3://<your-bucket>/jars/s3-tables-catalog-for-iceberg-runtime-0.1.5.jar",
        "--bigquery_connection_name": "Bigquery_connection",
        "--source_database": "cm_kawanago_tablebucket_namespace",
        "--source_table": "daily_sales",
        "--target_dataset": "cm_kawanago",
        "--target_table": "daily_sales_from_aws",
        "--gcp_project_id": "<your-gcp-project-id>"
    },
    "MaxRetries": 0,
    "Timeout": 60,
    "GlueVersion": "5.0",
    "MaxCapacity": 2.0,
    "Connections": {
        "Connections": ["Bigquery_connection"]
    }
}

"=== 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."