def get_results(self, query_execution_id: str) -> Iterator[Dict[str, Any]]:
"""
Get a query results and return a list of rows
:param query_execution_id: Query execution ID
:return: Iterator os lists
"""
res: Dict = self._client_athena.get_query_results(
QueryExecutionId=query_execution_id)
cols_info: List[Dict] = res["ResultSet"]["ResultSetMetadata"][
"ColumnInfo"]
athena_types: List[Tuple[str, str]] = [(x["Label"], x["Type"])
for x in cols_info]
logger.info(f"athena_types: {athena_types}")
python_types: List[Tuple[str, Optional[type]]] = [
(n, athena2python(dtype=t)) for n, t in athena_types
]
logger.info(f"python_types: {python_types}")
rows: List[Dict[str, List[Dict[str,
str]]]] = res["ResultSet"]["Rows"][1:]
for row in Athena._rows2row(rows=rows, python_types=python_types):
yield row
next_token: Optional[str] = res.get("NextToken")
while next_token is not None:
logger.info(f"next_token: {next_token}")
res = self._client_athena.get_query_results(
QueryExecutionId=query_execution_id, NextToken=next_token)
rows = res["ResultSet"]["Rows"]
for row in Athena._rows2row(rows=rows, python_types=python_types):
yield row
next_token = res.get("NextToken")
def get_table_python_types(self, database, table):
"""
Get all columns names and the related python types
:param database: Glue database's name
:param table: Glue table's name
:return: A dictionary as {"col name": "col python type"}
"""
dtypes = self.get_table_athena_types(database=database, table=table)
return {k: data_types.athena2python(v) for k, v in dtypes.items()}