Cloud Trail and Athena: Yet another solution design

SELECT
  account,
  useridentity.username,
  json_extract(responseelements, '$.assumedRoleUser.arn') as user_arn,
  count (useridentity.username) as count
FROM athena4cloudtrail
WHERE year = cast(date_format (current_timestamp - interval '1' day, '%Y') as int)
  AND month = cast(date_format (current_timestamp - interval '1' day, '%m') as int)
  AND day > cast(date_format (current_timestamp - interval '10' day, '%d') as int)
  AND eventname = 'AssumeRoleWithSAML'
  AND errorcode is null
GROUP BY account, useridentity.username, json_extract(responseelements, '$.assumedRoleUser.arn')
ORDER BY account, count DESC

from collections.abc import Generator
from datetime import datetime, timedelta

class PartitionGenerator(Generator):
    """ A generator class which generates batches of Partitions for cloudtrail organizational logs

        Parameters
        -------
        accounts : list(str)
        regions: list(str)
        startdate: datetime
        days_to_add: int
        page_size: int

        Returns
        -------
        partitions : list(account: str, region: str, year: int, month: int, day: int)
            Returns the next batch of length page_size for further processing

        Raises
        ------
        StopIteration

    """
    def __init__(self, accounts:list, regions:list, startdate: datetime, days_to_add: int, page_size: int):
        self.accounts = accounts
        self.regions = regions
        self.startdate = startdate
        self.enddate = startdate + timedelta(days=days_to_add)
        self.current = {
            'account': accounts[0],
            'region': regions[0],
            'date':startdate
        }
        self.pagesize = page_size
        self.accounts_processed = False

    def _update_current_partition(self, next_region_index, next_account_index) -> None:
        self.current['date']=self.startdate
        self.current = {
            'account': self.accounts[next_account_index],
            'region': self.regions[next_region_index],
            'date':self.startdate
        }
        return

    def send(self, arg):
        partition_values = []
        next_account_index = 0
        next_region_index = 0

        if self.accounts_processed:
            raise StopIteration

        for account_index, account in enumerate(self.accounts):
            if self.current['account'] != account: # Used in subsequent calls: Only proceed if we´ve reached the current account
                continue
            next_account_index = account_index
            for region_index, region in enumerate(self.regions):
                if self.current['region'] != region:  # Used in subsequent calls: Only proceed if we´ve reached the current region
                    continue
                while(self.current['date']<self.enddate):
                    partition_values.append({'Values': [
                        self.current['account'],
                        self.current['region'],
                        f"{self.current['date'].year}",
                        f"{self.current['date'].month:02d}",
                        f"{self.current['date'].day:02d}"
                        ]})
                    self.current['date']=self.current['date'] + timedelta(days=1)
                    if (self.current['date']>=self.enddate): # Datetime is out of range - move to next region and account (if necessary)
                        next_region_index = (region_index + 1) % len(self.regions)
                        if next_region_index == 0:
                            next_account_index = (account_index + 1) % len(self.accounts)
                            if next_account_index == 0:
                                self.accounts_processed = True
                        self._update_current_partition(
                            next_region_index,
                            next_account_index
                            )
                        if(len(partition_values) == self.pagesize):
                            return partition_values
                        break

                    if(len(partition_values) == self.pagesize):
                        return partition_values
            else:
                continue
            break

        self.accounts_processed = True
        if len(partition_values) > 0:
            return partition_values

    def throw(self, type, value=None, traceback=None):
        super().throw(type, value, traceback)

import sys
import os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../src')))
import boto3
import unittest
import random
from datetime import datetime, timedelta
from utils.partition_value_generator import PartitionGenerator


class TestPartitionValueGenerator(unittest.TestCase):
    @staticmethod
    def generate_example_partition_arguments() -> dict:
        # expected samples: 500
        return {
            'accounts':['123456789123','123456789124','123456789125','123456789126','123456789127'],
            'regions':['us-east-1','us-east-2','us-west-1','us-west-2','eu-west-1'],
            'startdate':datetime.today(),
            'days_to_add':20,
            'page_size':1
        }

    @staticmethod
    def get_random_partition_value_from_example() -> list:
        # returns one sample of the list of example partitions generated vai generate_example_partition_arguments
        account = random.choice(['123456789123','123456789124','123456789125','123456789126','123456789127'])
        region = random.choice(['us-east-1','us-east-2','us-west-1','us-west-2','eu-west-1'])
        date = datetime.today() + timedelta(days = random.choice(range(0,20)))
        return [
            account,
            region,
            str(date.year),
            f'{date.month:02d}',
            f'{date.day:02d}'
            ]

    @staticmethod
    def generate_random_partition_arguments() -> dict:
        all_regions = [
            'us-east-1',
            'us-east-2',
            'us-west-1',
            'us-west-2',
            'eu-west-1',
            'eu-central-1',
            'ap-northeast-1',
            'ap-northeast-2',
            'ap-southeast-1',
            'ap-southeast-2',
            'ap-south-1',
            'sa-east-1'
            ]
        regions = random.sample(all_regions, random.randint(1,len(all_regions)))
        accounts = []
        for _ in range(random.randint(1,20)):
            random_account_id = ''.join(str(random.randint(0,9)) for _ in range(12))
            accounts.append(random_account_id)
        startdate = datetime.today()
        days_to_add = random.randint(1,10)
        page_size = 1
        return {
            'accounts':accounts,
            'regions':regions,
            'startdate':startdate,
            'days_to_add':days_to_add,
            'page_size':page_size
        }

    def test_partition_value_count(self) -> None:
        arguments = self.generate_random_partition_arguments()
        partition_generator = PartitionGenerator(**arguments)
        expected_count = len(arguments['accounts']) * len(arguments['regions']) * arguments['days_to_add']

        count = len([partition for partition in partition_generator])

        self.assertEqual(
            expected_count,
            count,
            f'Partitioncount {count} doesn´t match for random generated partion request {arguments}'
            )

    def test_partition_pagesize(self) -> None:
        arguments = self.generate_example_partition_arguments()
        arguments['page_size'] = random.choice([5,10,25,50,100])
        partition_generator = PartitionGenerator(**arguments)

        count = 0
        for partition_batch in partition_generator:
            self.assertEqual(
                arguments['page_size'],
                len(partition_batch),
                f'Partition pagesize {len(partition_batch)} is not correct'
            )
            count = count + arguments['page_size']
        self.assertEqual(
            500,
            count,
            f'Partitioncount {count} doesn´t match for random generated partion request {arguments}'
            )

    def test_partition_content(self) -> None:
        arguments = self.generate_example_partition_arguments()
        arguments['page_size'] = 500
        partition_generator = PartitionGenerator(**arguments)
        partitions = []

        for partition_batch in partition_generator:
            partitions.extend(partition_batch)

        for _ in range(10):
            random_partition = self.get_random_partition_value_from_example()
            partition_filtered = next((partition for partition in partitions if partition['Values']==random_partition), None)
            self.assertIsNotNone(
                partition_filtered,
                f'Partion {random_partition} not found in result.')

unittest.main()