Athenaの実行結果をダウンロードして読み込むモジュール

from lib import athena
import csv

# 一時ファイルの保存先のs3のパスを指定
client = athena.AthenaClient('s3://ninomiyt-example-bucket/athena/')

# SQLを実行し、結果のファイルをダウンロードする
job = client.execute('SELECT * FROM table_name', 'database_name')
filepath = job.fetch('/tmp/athena/', sleep_sec=5)

# ダウンロードしたcsvの読み込み
with filepath.open('r') as f:
    reader = csv.DictReader(f)
    for row in reader:
        pass
        # 行ごとの処理

import csv
from pathlib import Path
from time import sleep
import boto3
from retry import retry
from . import s3


REGION_NAME = 'ap-northeast-1'


class AthenaClient(object):
    """Athenaのクライアント"""

    def __init__(self, output_location):
        """Args:
            output_location (str): 出力先を表す"s3://{バケット名}/{パス}"の文字列

        """
        self.athena = boto3.client('athena', region_name=REGION_NAME)
        self.output_location = output_location

    @retry(tries=4, delay=5)
    def execute(self, query_string, database):
        """Athenaに対してSQLを実行する。

        Args:
            query_string (str): 実行するSQL
            database (str): 実行対象のDB名
        Returns:
            QueryExecutionJob: クエリの実行結果のオブジェクト

        """
        res = self.athena.start_query_execution(
            QueryString=query_string,
            QueryExecutionContext={'Database': database},
            ResultConfiguration={'OutputLocation': self.output_location}
        )
        return QueryExecutionJob(self.athena, res)


class QueryExecutionJob(object):
    """各クエリの実行結果を表すクラス"""

    def __init__(self, athena, query_exec):
        """Args:
            athena (boto3.client): クエリを実行したboto3のクライアント
            query_exec (Dict): start_query_executionのAPIの戻り値

        """
        self.athena = athena
        self.query_exec_id = query_exec['QueryExecutionId']

    def fetch(self, localdir, *, sleep_sec):
        """実行結果をローカルにダウンロードする。

        Args:
            localdir (str/pathlib.Path): 保存先のディレクトリ
            sleep_sec (int/float): クエリが完了していない場合に待つ秒数
        Returns:
            pathlib.Path: ローカルにダウンロードしたテキストファイルの情報
        Raises:
            QueryExecutionError: クエリが失敗した場合
            通信失敗時のエラー

        """
        while not self._is_complete():
            sleep(sleep_sec)
        return self._download_csv_file(localdir)

    def _is_complete(self):
        """クエリが完了しているか調べる

        Returns:
            bool: クエリが完了していればTrue
        Raises:
            AthenaExecutionError: クエリが失敗した場合
            通信失敗時のエラー

        """
        state = self._get_query_execution_state()
        if state in ('FAILED', 'CANCELLED'):
            raise QueryExecutionError(
                f'{self.query_exec_id} のジョブがState: {state} で失敗しました')
        return state == 'SUCCEEDED'

    def _download_csv_file(self, localdir):
        """S3の実行結果のファイルをダウンロードする
        保存先のディレクトリが存在しない場合は作成する

        Args:
            localdir(str/pathlib.Path): 保存先のディレクトリ
        Returns:
            pathlib.Path: ローカルにダウンロードしたテキストファイルの情報
        Raises:
            通信失敗時のエラー

        """
        s3_location = self._get_output_location()
        localpath = Path(localdir) / s3_location.split('/')[-1]
        # 保存先を/tmp/以下で、元々ディレクトリが存在しないことも想定しています
        localpath.parent.mkdir(parents=True, exist_ok=True)
        s3.download_file(s3_location, localpath)
        return localpath

    @retry(tries=4, delay=5)
    def _get_query_execution_state(self):
        """クエリの実行状況を返す。

        Returns:
            str: クエリの実行状況 (RUNNING|SUCCEEDED|FAILED|CANCELLED)

        """
        res = self.athena.get_query_execution(
            QueryExecutionId=self.query_exec_id)
        return res['QueryExecution']['Status']['State']

    @retry(tries=4, delay=5)
    def _get_output_location(self):
        """アウトプットCSVファイルの場所を取得する。

        Returns:
            str: s3://{バケット名}/{パス} 形式の文字列
        """
        res = self.athena.get_query_execution(
            QueryExecutionId=self.query_exec_id)
        return res['QueryExecution']['ResultConfiguration']['OutputLocation']


class QueryExecutionError(Exception):
    """クエリが失敗・中止されたことを表すエラー"""
    pass

import csv
from urllib.parse import urlparse
import boto3
from retry import retry


@retry(tries=4, delay=5)
def download_file(s3_path, localpath):
    """S3のファイルをローカルにダウンロードする。

    Args:
        s3_path (str): S3にあるファイルのパス
        localpath (str/pathlib.Path): ローカルのファイルパス

    """
    s3 = boto3.resource('s3')
    bk, key = _split_s3_path(s3_path)
    bucket = s3.Bucket(bk)
    bucket.download_file(key, str(localpath))


def _split_s3_path(s3_path):
    """s3://{バケット名}/{パス}を分割する。
    こちらを参考にしました。
    https://stackoverflow.com/questions/42641315/s3-urls-get-bucket-name-and-path

    Args:
        s3_path(str): s3://{バケット名}/{パス}の文字列
    Returns:
        tuple: 以下の値を返す
            - str: バケット名
            - str: S3上のパス

    """
    parsed = urlparse(s3_path)
    return parsed.netloc, parsed.path.strip('/')