def extract_athena_table_columns(
database: str, table: str,
boto3_session: boto3.Session) -> List[Dict[str, str]]:
"""Extract athena columns data types from table and raising an exception if not exist."""
dtypes: Optional[Dict[str, str]] = catalog.get_table_types(
database=database, table=table, boto3_session=boto3_session)
if dtypes is None:
raise exceptions.InvalidArgument(
f"{database}.{table} does not exist on Athena.")
return [{
"Name": name,
"Type": _data_types.athena2quicksight(dtype=dtype)
} for name, dtype in dtypes.items()]
def extract_athena_query_columns(
sql: str, data_source_arn: str, account_id: str,
boto3_session: boto3.Session) -> List[Dict[str, str]]:
"""Extract athena columns data types from a SQL query."""
data_sources: List[Dict[str, Any]] = list_data_sources(
account_id=account_id, boto3_session=boto3_session)
data_source: Dict[str, Any] = [
x for x in data_sources if x["Arn"] == data_source_arn
][0]
workgroup: str = data_source["DataSourceParameters"]["AthenaParameters"][
"WorkGroup"]
sql_wrapped: str = f"/* QuickSight */\nSELECT ds.* FROM ( {sql} ) ds LIMIT 0"
query_id: str = athena.start_query_execution(sql=sql_wrapped,
workgroup=workgroup,
boto3_session=boto3_session)
athena.wait_query(query_execution_id=query_id, boto3_session=boto3_session)
dtypes: Dict[str, str] = athena.get_query_columns_types(
query_execution_id=query_id, boto3_session=boto3_session)
return [{
"Name": name,
"Type": _data_types.athena2quicksight(dtype=dtype)
} for name, dtype in dtypes.items()]