support for pagination for timestream.list_databases and timestream.list_tables (#1846)

Co-authored-by: kukushking <[email protected]>
Co-authored-by: Lucas Hanson <[email protected]>

Author: 3 people

awswrangler/catalog/_get.py

@@ -242,7 +242,7 @@ def get_tables(
 
     Note
     ----
-    Please, does not filter using name_contains and name_prefix/name_suffix at the same time.
+    Please, do not filter using name_contains and name_prefix/name_suffix at the same time.
     Only name_prefix and name_suffix can be combined together.
 
     Parameters
@@ -279,7 +279,7 @@ def get_tables(
     args: Dict[str, str] = {}
     if (name_prefix is not None) and (name_suffix is not None) and (name_contains is not None):
         raise exceptions.InvalidArgumentCombination(
-            "Please, does not filter using name_contains and "
+            "Please, do not filter using name_contains and "
             "name_prefix/name_suffix at the same time. Only "
             "name_prefix and name_suffix can be combined together."
         )

awswrangler/timestream.py

@@ -536,8 +536,14 @@ def list_databases(
 
     """
     client: boto3.client = _utils.client(service_name="timestream-write", session=boto3_session)
-    dbs = client.list_databases()
-    return [db["DatabaseName"] for db in dbs["Databases"]]
+
+    response: Dict[str, Any] = client.list_databases()
+    dbs: List[str] = [db["DatabaseName"] for db in response["Databases"]]
+    while "nextToken" in response:
+        response = client.list_databases(nextToken=response["nextToken"])
+        dbs += [db["DatabaseName"] for db in response["Databases"]]
+
+    return dbs
 
 
 def list_tables(database: Optional[str] = None, boto3_session: Optional[boto3.Session] = None) -> List[str]:
@@ -573,9 +579,11 @@ def list_tables(database: Optional[str] = None, boto3_session: Optional[boto3.Se
 
     """
     client: boto3.client = _utils.client(service_name="timestream-write", session=boto3_session)
-    if database:
-        tables = client.list_tables(DatabaseName=database)
-    else:
-        tables = client.list_tables()
-
-    return [tbl["TableName"] for tbl in tables["Tables"]]
+    args = {} if database is None else {"DatabaseName": database}
+    response: Dict[str, Any] = client.list_tables(**args)
+    tables: List[str] = [tbl["TableName"] for tbl in response["Tables"]]
+    while "nextToken" in response:
+        response = client.list_tables(**args, nextToken=response["nextToken"])
+        tables += [tbl["TableName"] for tbl in response["Tables"]]
+
+    return tables