import logging

from datetime import datetime
from psi_elt.abstract_elt_job import AbstractELTJob, LOADER_ELT_JOB_TYPE
from typeguard import typechecked
from typing import Any


@typechecked
class CustomLoader(AbstractELTJob):
    '''A CustomLoader is the second of the ELT jobs that is responsible for
    loading raw data into S3 in a queryable format.

    This class should be treated as abstract and should be extended.
    Child classes must override the _load function.
    '''

    def __init__(
        self,
        name: str,
        capture_time: datetime,
        aws_session: Any = None,
        logger: logging.Logger = None
    ):
        super().__init__(
            name=name,
            elt_job_type=LOADER_ELT_JOB_TYPE,
            capture_time=capture_time,
            aws_session=aws_session,
            logger=logger
        )

    def load(self, is_test: bool, **load_kwargs):
        load_kwargs = load_kwargs or {}
        load_kwargs = {k: v for k, v in load_kwargs.items() if v is not None}
        data = self._load(is_test, **load_kwargs)
        return data

    def _load(self, is_test: bool, **load_kwargs) -> Any:
        raise Exception('CustomLoader._load must be overridden by the child class')
