Python×Aurora DSQLで作るリアルタイムETLパイプライン（コード付き）

# 仮想環境の作成と有効化
python3 -m venv aurora_dsql_etl
source aurora_dsql_etl/bin/activate

# 必要ライブラリのインストール
pip install boto3 botocore
pip install psycopg[binary]
pip install sqlalchemy
pip install pandas
pip install requests

import boto3
import pandas as pd
from sqlalchemy import create_engine, text
from sqlalchemy.engine import URL
import logging
import time
from typing import Optional

class AuroraDSQLConnector:
    """Aurora DSQL接続管理クラス"""
    
    def __init__(self, hostname: str, region: str = 'us-east-1'):
        self.hostname = hostname
        self.region = region
        self.dsql_client = boto3.client('dsql', region_name=region)
        self.engine = None
        self.logger = logging.getLogger(__name__)
        
    def create_engine(self) -> None:
        """SQLAlchemyエンジンを作成"""
        try:
            # IAM認証トークン生成（1時間有効）
            password_token = self.dsql_client.generate_db_connect_admin_auth_token(
                hostname=self.hostname,
                region=self.region,
                expires_in=3600
            )
            
            # PostgreSQL互換接続URL作成
            url = URL.create(
                "postgresql",
                username="admin",
                password=password_token,
                host=self.hostname,
                database="postgres"
            )
            
            # SSL必須設定でエンジン作成
            self.engine = create_engine(
                url, 
                connect_args={"sslmode": "require"},
                pool_size=10,
                max_overflow=20,
                pool_pre_ping=True
            )
            
            self.logger.info(f"Aurora DSQLエンジン作成完了: {self.hostname}")
            
        except Exception as e:
            self.logger.error(f"エンジン作成エラー: {e}")
            raise

import uuid
from typing import List, Dict, Any

class DataTransformer:
    """データ変換クラス"""
    
    def __init__(self):
        self.logger = logging.getLogger(__name__)
        
    def clean_and_validate(self, df: pd.DataFrame) -> pd.DataFrame:
        """データクリーニングとバリデーション"""
        try:
            # 元のデータ件数を記録
            original_count = len(df)
            
            # 重複行の削除
            df = df.drop_duplicates()
            
            # 空値の処理
            df = df.dropna(subset=['id'])  # 必須フィールドの空値は除外
            df = df.fillna('')  # その他は空文字で埋める
            
            # データ型の正規化
            df = self._normalize_data_types(df)
            
            # Aurora DSQL用にUUID主キー追加（SERIAL非対応のため）
            if 'uuid_id' not in df.columns:
                df['uuid_id'] = [str(uuid.uuid4()) for _ in range(len(df))]
            
            cleaned_count = len(df)
            self.logger.info(f"データクリーニング完了: {original_count} → {cleaned_count}件")
            
            return df
            
        except Exception as e:
            self.logger.error(f"データクリーニングエラー: {e}")
            raise
            
    def _normalize_data_types(self, df: pd.DataFrame) -> pd.DataFrame:
        """データ型の正規化"""
        # 日時フィールドの正規化
        for col in df.columns:
            if 'date' in col.lower() or 'time' in col.lower():
                try:
                    df[col] = pd.to_datetime(df[col], errors='coerce')
                except:
                    pass
                    
        # 数値フィールドの正規化
        for col in df.select_dtypes(include=['object']).columns:
            # 数値らしい文字列を数値に変換
            try:
                numeric_series = pd.to_numeric(df[col], errors='coerce')
                if not numeric_series.isna().all():
                    df[col] = numeric_series
            except:
                pass
                
        return df
        
    def apply_business_rules(self, df: pd.DataFrame, rules: List[Dict[str, Any]]) -> pd.DataFrame:
        """ビジネスルールの適用"""
        try:
            for rule in rules:
                rule_type = rule.get('type')
                
                if rule_type == 'filter':
                    # フィルタリングルール
                    condition = rule.get('condition')
                    df = df.query(condition)
                    
                elif rule_type == 'transform':
                    # 変換ルール
                    column = rule.get('column')
                    function = rule.get('function')
                    if hasattr(df[column], function):
                        df[column] = getattr(df[column], function)()
                        
                elif rule_type == 'enrich':
                    # エンリッチメントルール
                    new_column = rule.get('new_column')
                    logic = rule.get('logic')
                    df[new_column] = df.eval(logic)
                    
            self.logger.info(f"ビジネスルール適用完了: {len(rules)}ルール処理")
            return df
            
        except Exception as e:
            self.logger.error(f"ビジネスルール適用エラー: {e}")
            raise

class DataLoader:
    """データ格納クラス"""
    
    def __init__(self, db_connector: AuroraDSQLConnector):
        self.db = db_connector
        self.logger = logging.getLogger(__name__)
        
    def create_table_if_not_exists(self, table_name: str, df: pd.DataFrame) -> None:
        """テーブルが存在しない場合は作成"""
        try:
            # DataFrameのカラム情報からCREATE TABLE文を生成
            columns = []
            for col_name, dtype in df.dtypes.items():
                if col_name == 'uuid_id':
                    columns.append(f"{col_name} UUID PRIMARY KEY")
                elif 'datetime' in str(dtype):
                    columns.append(f"{col_name} TIMESTAMP")
                elif 'int' in str(dtype):
                    columns.append(f"{col_name} INTEGER")
                elif 'float' in str(dtype):
                    columns.append(f"{col_name} DECIMAL")
                else:
                    columns.append(f"{col_name} TEXT")
                    
            columns_sql = ",\n    ".join(columns)
            
            create_table_sql = f"""
            CREATE TABLE IF NOT EXISTS {table_name} (
                {columns_sql},
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )
            """
            
            self.db.execute_with_retry(create_table_sql)
            self.logger.info(f"テーブル作成/確認完了: {table_name}")
            
        except Exception as e:
            self.logger.error(f"テーブル作成エラー: {e}")
            raise
            
    def bulk_insert(self, table_name: str, df: pd.DataFrame, chunk_size: int = 1000) -> None:
        """バルクインサート（バッチ処理）"""
        try:
            total_rows = len(df)
            inserted_rows = 0
            
            # チャンクサイズごとに分割して処理
            for chunk_start in range(0, total_rows, chunk_size):
                chunk_end = min(chunk_start + chunk_size, total_rows)
                chunk_df = df.iloc[chunk_start:chunk_end]
                
                # プリペアドステートメント用のINSERT文生成
                columns = list(chunk_df.columns)
                placeholders = ", ".join([f":{col}" for col in columns])
                
                insert_sql = f"""
                INSERT INTO {table_name} ({", ".join(columns)})
                VALUES ({placeholders})
                """
                
                # データをdict形式に変換
                data_dicts = chunk_df.to_dict('records')
                
                # バッチ実行
                with self.db.engine.connect() as conn:
                    conn.execute(text(insert_sql), data_dicts)
                    conn.commit()
                    
                inserted_rows += len(chunk_df)
                self.logger.info(f"バッチ挿入進捗: {inserted_rows}/{total_rows}件")
                
            self.logger.info(f"バルクインサート完了: {table_name}に{total_rows}件挿入")
            
        except Exception as e:
            self.logger.error(f"バルクインサートエラー: {e}")
            raise
            
    def upsert_data(self, table_name: str, df: pd.DataFrame, key_columns: List[str]) -> None:
        """UPSERT処理（INSERT ON CONFLICT）"""
        try:
            columns = list(df.columns)
            data_dicts = df.to_dict('records')
            
            # ON CONFLICT UPDATE用のSET句生成
            update_columns = [col for col in columns if col not in key_columns]
            set_clause = ", ".join([f"{col} = EXCLUDED.{col}" for col in update_columns])
            
            conflict_columns = ", ".join(key_columns)
            
            upsert_sql = f"""
            INSERT INTO {table_name} ({", ".join(columns)})
            VALUES ({", ".join([f":{col}" for col in columns])})
            ON CONFLICT ({conflict_columns})
            DO UPDATE SET {set_clause}, updated_at = CURRENT_TIMESTAMP
            """
            
            # バッチ実行
            with self.db.engine.connect() as conn:
                conn.execute(text(upsert_sql), data_dicts)
                conn.commit()
                
            self.logger.info(f"UPSERT完了: {table_name}に{len(df)}件処理")
            
        except Exception as e:
            self.logger.error(f"UPSERTエラー: {e}")
            raise

(self, table_name: str, df: pd.DataFrame, key_columns: List[str]) -> None:
        """UPSERT処理（INSERT ON CONFLICT）"""
        try:
            columns = list(df.columns)
            data_dicts = df.to_dict('records')
            
            # ON CONFLICT UPDATE用のSET句生成
            update_columns = [col for col in columns if col not in key_columns]
            set_clause = ", ".join([f"{col} = EXCLUDED.{col}" for col in update_columns])
            
            conflict_columns = ", ".join(key_columns)
            
            upsert_sql = f"""
            INSERT INTO {table_name} ({", ".join(columns)})
            VALUES ({", ".join([f":{col}" for col in columns])})
            ON CONFLICT ({conflict_columns})
            DO UPDATE SET {set_clause}, updated_at = CURRENT_TIMESTAMP
            """
            
            # バッチ実行
            with self.db.engine.connect() as conn:
                conn.execute(text(upsert_sql), data_dicts)
                conn.commit()
                
            self.logger.info(f"UPSERT完了: {table_name}に{len(df)}件処理")
            
        except Exception as e:
            self.logger.error(f"UPSERTエラー: {e}")
            raise

class AuroraDSQLETLPipeline:
    """Aurora DSQL ETLパイプライン統合クラス"""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.logger = self._setup_logger()
        
        # 各コンポーネント初期化
        self.db_connector = AuroraDSQLConnector(
            hostname=config['aurora_dsql']['hostname'],
            region=config['aurora_dsql']['region']
        )
        self.extractor = DataExtractor(config['data_source']['base_url'])
        self.transformer = DataTransformer()
        self.loader = DataLoader(self.db_connector)
        
    def _setup_logger(self) -> logging.Logger:
        """ログ設定"""
        logger = logging.getLogger('etl_pipeline')
        logger.setLevel(logging.INFO)
        
        if not logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            logger.addHandler(handler)
            
        return logger
        
    def run_pipeline(self, pipeline_name: str) -> Dict[str, Any]:
        """ETLパイプライン実行"""
        start_time = datetime.now()
        pipeline_config = self.config['pipelines'][pipeline_name]
        
        try:
            self.logger.info(f"ETLパイプライン開始: {pipeline_name}")
            
            # 1. Extract（抽出）
            self.logger.info("データ抽出フェーズ開始")
            if pipeline_config['source_type'] == 'api':
                raw_data = self.extractor.extract_api_data(
                    endpoint=pipeline_config['endpoint'],
                    params=pipeline_config.get('params', {})
                )
            elif pipeline_config['source_type'] == 'file':
                raw_data = self.extractor.extract_file_data(
                    file_path=pipeline_config['file_path']
                )
            else:
                raise ValueError(f"未対応のソースタイプ: {pipeline_config['source_type']}")
                
            # 2. Transform（変換）
            self.logger.info("データ変換フェーズ開始")
            cleaned_data = self.transformer.clean_and_validate(raw_data)
            
            if 'business_rules' in pipeline_config:
                transformed_data = self.transformer.apply_business_rules(
                    cleaned_data, 
                    pipeline_config['business_rules']
                )
            else:
                transformed_data = cleaned_data
                
            # 3. Load（格納）
            self.logger.info("データ格納フェーズ開始")
            
            # Aurora DSQL接続確立
            if not self.db_connector.engine:
                self.db_connector.create_engine()
                
            target_table = pipeline_config['target_table']
            
            # テーブル作成（存在しない場合）
            self.loader.create_table_if_not_exists(target_table, transformed_data)
            
            # データ格納方式の選択
            load_mode = pipeline_config.get('load_mode', 'insert')
            if load_mode == 'upsert':
                self.loader.upsert_data(
                    target_table, 
                    transformed_data, 
                    pipeline_config['key_columns']
                )
            else:
                self.loader.bulk_insert(target_table, transformed_data)
                
            # 実行結果サマリー
            end_time = datetime.now()
            execution_time = (end_time - start_time).total_seconds()
            
            result = {
                'pipeline_name': pipeline_name,
                'status': 'success',
                'start_time': start_time,
                'end_time': end_time,
                'execution_time_seconds': execution_time,
                'processed_records': len(transformed_data),
                'target_table': target_table
            }
            
            self.logger.info(f"ETLパイプライン完了: {pipeline_name} ({execution_time:.2f}秒)")
            return result
            
        except Exception as e:
            end_time = datetime.now()
            execution_time = (end_time - start_time).total_seconds()
            
            error_result = {
                'pipeline_name': pipeline_name,
                'status': 'failed',
                'start_time': start_time,
                'end_time': end_time,
                'execution_time_seconds': execution_time,
                'error_message': str(e)
            }
            
            self.logger.error(f"ETLパイプライン失敗: {pipeline_name} - {e}")
            return error_result
            
    def run_scheduled_pipeline(self, schedule_config: Dict[str, Any]) -> None:
        """スケジュール実行（AWS Lambdaやcronと組み合わせて使用）"""
        for pipeline_name in schedule_config['pipelines']:
            try:
                result = self.run_pipeline(pipeline_name)
                
                # 実行結果をログテーブルに記録
                self._log_execution_result(result)
                
                # 失敗時の通知処理
                if result['status'] == 'failed':
                    self._send_failure_notification(result)
                    
            except Exception as e:
                self.logger.error(f"スケジュール実行エラー: {pipeline_name} - {e}")
                
    def _log_execution_result(self, result: Dict[str, Any]) -> None:
        """実行結果ログの記録"""
        try:
            log_data = pd.DataFrame([result])
            self.loader.bulk_insert('etl_execution_log', log_data)
        except Exception as e:
            self.logger.error(f"実行ログ記録エラー: {e}")
            
    def _send_failure_notification(self, result: Dict[str, Any]) -> None:
        """失敗通知の送信（SNS等と連携）"""
        # 実装例：AWS SNSによる通知
        try:
            # SNS通知ロジックをここに実装
            self.logger.warning(f"ETLパイプライン失敗通知: {result['pipeline_name']}")
        except Exception as e:
            self.logger.error(f"失敗通知送信エラー: {e}")

# config.json
{
    "aurora_dsql": {
        "hostname": "your-cluster.dsql.us-east-1.on.aws",
        "region": "us-east-1"
    },
    "data_source": {
        "base_url": "https://api.example.com"
    },
    "pipelines": {
        "user_data_pipeline": {
            "source_type": "api",
            "endpoint": "users",
            "params": {
                "limit": 1000,
                "format": "json"
            },
            "target_table": "users",
            "load_mode": "upsert",
            "key_columns": ["user_id"],
            "business_rules": [
                {
                    "type": "filter",
                    "condition": "status == 'active'"
                },
                {
                    "type": "enrich",
                    "new_column": "full_name",
                    "logic": "first_name + ' ' + last_name"
                }
            ]
        },
        "sales_data_pipeline": {
            "source_type": "file",
            "file_path": "/data/sales.csv",
            "target_table": "sales",
            "load_mode": "insert"
        }
    }
}

# main.py
import json
import sys
from etl_pipeline import AuroraDSQLETLPipeline

def main():
    """メイン処理"""
    try:
        # 設定ファイル読み込み
        with open('config.json', 'r', encoding='utf-8') as f:
            config = json.load(f)
            
        # ETLパイプライン初期化
        pipeline = AuroraDSQLETLPipeline(config)
        
        # コマンドライン引数で実行パイプラインを指定
        if len(sys.argv) > 1:
            pipeline_name = sys.argv[1]
            
            if pipeline_name in config['pipelines']:
                result = pipeline.run_pipeline(pipeline_name)
                print(f"パイプライン実行結果: {result}")
            else:
                print(f"エラー: 未定義のパイプライン '{pipeline_name}'")
                sys.exit(1)
        else:
            # 全パイプライン実行
            for pipeline_name in config['pipelines'].keys():
                result = pipeline.run_pipeline(pipeline_name)
                print(f"パイプライン実行結果: {result}")
                
    except Exception as e:
        print(f"エラー: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()

# 実行例
# python main.py user_data_pipeline  # 特定パイプライン実行
# python main.py                     # 全パイプライン実行