timestream.query: add QueryId and NextToken to df attributes

awswrangler/timestream.py

@@ -122,8 +122,12 @@ def _process_row(schema: List[Dict[str, str]], row: Dict[str, Any]) -> List[Any]
     return row_processed
 
 
-def _rows_to_df(rows: List[List[Any]], schema: List[Dict[str, str]]) -> pd.DataFrame:
+def _rows_to_df(
+    rows: List[List[Any]], schema: List[Dict[str, str]], df_metadata: Optional[Dict[str, str]] = None
+) -> pd.DataFrame:
     df = pd.DataFrame(data=rows, columns=[c["name"] for c in schema])
+    if df_metadata:
+        df.attrs = df_metadata
     for col in schema:
         if col["type"] == "VARCHAR":
             df[col["name"]] = df[col["name"]].astype("string")
@@ -143,7 +147,7 @@ def _process_schema(page: Dict[str, Any]) -> List[Dict[str, str]]:
 
 
 def _paginate_query(
-    sql: str, pagination_config: Optional[Dict[str, Any]], boto3_session: Optional[boto3.Session] = None
+    sql: str, chunked: bool, pagination_config: Optional[Dict[str, Any]], boto3_session: Optional[boto3.Session] = None
 ) -> Iterator[pd.DataFrame]:
     client: boto3.client = _utils.client(
         service_name="timestream-query",
@@ -161,7 +165,13 @@ def _paginate_query(
         for row in page["Rows"]:
             rows.append(_process_row(schema=schema, row=row))
         if len(rows) > 0:
-            yield _rows_to_df(rows, schema)
+            df_metadata = {}
+            if chunked:
+                if "NextToken" in page:
+                    df_metadata["NextToken"] = page["NextToken"]
+                df_metadata["QueryId"] = page["QueryId"]
+
+            yield _rows_to_df(rows, schema, df_metadata)
         rows = []
 
 
@@ -289,9 +299,10 @@ def query(
     >>> df = wr.timestream.query('SELECT * FROM "sampleDB"."sampleTable" ORDER BY time DESC LIMIT 10')
 
     """
-    result_iterator = _paginate_query(sql, pagination_config, boto3_session)
+    result_iterator = _paginate_query(sql, chunked, pagination_config, boto3_session)
     if chunked:
         return result_iterator
+
     # Prepending an empty DataFrame ensures returning an empty DataFrame if result_iterator is empty
     return pd.concat(itertools.chain([pd.DataFrame()], result_iterator), ignore_index=True)
 

tests/test_timestream.py

@@ -47,6 +47,7 @@ def test_basic_scenario(timestream_database_and_table, pagination):
         pagination_config=pagination,
     )
     assert df.shape == (3, 8)
+    assert df.attrs == {}
 
 
 @pytest.mark.parametrize("chunked", [False, True])
@@ -114,6 +115,7 @@ def test_chunked_scenario(timestream_database_and_table):
         ),
         shapes,
     ):
+        assert "QueryId" in df.attrs
         assert df.shape == shape