athena client type stubs

athena_client.py

from mypy_extensions import TypedDict

QueryExecutionContext = TypedDict('QueryExecutionContext', {
    'Database': str,
})
EncryptionConfiguration = TypedDict('EncryptionConfiguration', {
    'EncryptionOption': str,
    'KmsKey': str,
}, total=False)
ResultConfiguration = TypedDict('ResultConfiguration', {
    'EncryptionConfiguration': EncryptionConfiguration,
    'OutputLocation': str,
})
Statistics = TypedDict('Statistics', {
    'DataScannedInBytes': int,
    'EngineExecutionTimeInMillis': int,
})
Status = TypedDict('Status', {
    'CompletionDateTime': datetime.datetime,
    'State': str,
    'StateChangeReason': str,
    'SubmissionDateTime': datetime.datetime,
})
QueryExecution = TypedDict('QueryExecution', {
    'Query': str,
    'QueryExecutionContext': QueryExecutionContext,
    'QueryExecutionId': str,
    'ResultConfiguration': ResultConfiguration,
    'Statistics': Statistics,
    'Status': Status,
})
GetQueryExecutionResponse = TypedDict('GetQueryExecutionResponse', {
    'QueryExecution': QueryExecution,
})

StartQueryExecutionResponse = TypedDict('StartQueryExecutionResponse', {
    'QueryExecutionId': str
})
ListQueryExecutionsResponse = TypedDict('ListQueryExecutionsResponse', {
    'NextToken': Optional[str],
    'QueryExecutionIds': List[str]
})


class EncryptionOption(enum.Enum):
    SSE_S3 = 'SSE_S3'
    SSE_KMS = 'SSE_KMS'
    CSE_KMS = 'CSE_KMS'


class AthenaClient:
    def get_query_execution(self, *, QueryExecutionId: str) -> GetQueryExecutionResponse:
        pass

    def stop_query_execution(self, *, QueryExecutionId: str):
        pass

    def list_query_executions(
            self, *,
            MaxResults: Optional[int]=None) -> ListQueryExecutionsResponse:
        pass

    def start_query_execution(
            self, *,
            # 28 - 128 characters
            ClientRequestToken: str='',
            QueryExecutionContext: QueryExecutionContext=None,
            QueryString: str,
            ResultConfiguration: ResultConfiguration) -> StartQueryExecutionResponse:
        pass