enable list timestream databases and tables (#1345)

* extended timestream.py with two functions to list databases and list tables

* added tests for timestream list db and tbl functions

* blackened

* more extensive tests
less assumptions (removed assumed count of dbs of 1)

* timestream.py: avoided None error

Co-authored-by: Thomas Niebler <[email protected]>
Co-authored-by: jaidisido <[email protected]>

Author: 3 people

awswrangler/timestream.py

@@ -482,3 +482,75 @@ def delete_table(
     """
     client: boto3.client = _utils.client(service_name="timestream-write", session=boto3_session)
     client.delete_table(DatabaseName=database, TableName=table)
+
+
+def list_databases(
+    boto3_session: Optional[boto3.Session] = None,
+) -> List[str]:
+    """
+    List all databases in timestream.
+
+    Parameters
+    ----------
+    boto3_session : boto3.Session(), optional
+        Boto3 Session. The default boto3 Session will be used if boto3_session receive None.
+
+    Returns
+    -------
+    List[str]
+        a list of available timestream databases.
+
+    Examples
+    --------
+    Querying the list of all available databases
+
+    >>> import awswrangler as wr
+    >>> wr.timestream.list_databases()
+    ... ["database1", "database2"]
+
+
+    """
+    client: boto3.client = _utils.client(service_name="timestream-write", session=boto3_session)
+    dbs = client.list_databases()
+    return [db["DatabaseName"] for db in dbs["Databases"]]
+
+
+def list_tables(database: Optional[str] = None, boto3_session: Optional[boto3.Session] = None) -> List[str]:
+    """
+    List tables in timestream.
+
+    Parameters
+    ----------
+    database: str
+        Database name. If None, all tables in Timestream will be returned. Otherwise, only the tables inside the
+        given database are returned.
+    boto3_session : boto3.Session(), optional
+        Boto3 Session. The default boto3 Session will be used if boto3_session receive None.
+
+    Returns
+    -------
+    List[str]
+        A list of table names.
+
+    Examples
+    --------
+    Listing all tables in timestream across databases
+
+    >>> import awswrangler as wr
+    >>> wr.timestream.list_tables()
+    ... ["table1", "table2"]
+
+    Listing all tables in timestream in a specific database
+
+    >>> import awswrangler as wr
+    >>> wr.timestream.list_tables(DatabaseName="database1")
+    ... ["table1"]
+
+    """
+    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"]]

tests/test_timestream.py

@@ -211,3 +211,50 @@ def test_multimeasure_scenario(timestream_database_and_table):
         """,
     )
     assert df.shape == (3, 6)
+
+
+def test_list_databases(timestream_database_and_table):
+    dbs = wr.timestream.list_databases()
+
+    assert timestream_database_and_table in dbs
+    dummy_db_name = f"{timestream_database_and_table}_2"
+
+    wr.timestream.create_database(dummy_db_name)
+    dbs_tmp = wr.timestream.list_databases()
+
+    assert timestream_database_and_table in dbs_tmp
+    assert dummy_db_name in dbs_tmp
+    assert len(dbs_tmp) == len(dbs) + 1
+
+    wr.timestream.delete_database(dummy_db_name)
+
+    dbs_tmp = wr.timestream.list_databases()
+    assert timestream_database_and_table in dbs_tmp
+    assert dummy_db_name not in dbs_tmp
+    assert len(dbs_tmp) == len(dbs)
+    assert dbs_tmp == dbs
+
+
+def test_list_tables(timestream_database_and_table):
+    all_tables = wr.timestream.list_tables()
+
+    assert timestream_database_and_table in all_tables
+
+    tables_in_db = wr.timestream.list_tables(database=timestream_database_and_table)
+    assert timestream_database_and_table in tables_in_db
+    assert len(tables_in_db) <= len(all_tables)
+
+    wr.timestream.create_table(
+        database=timestream_database_and_table,
+        table=f"{timestream_database_and_table}_2",
+        memory_retention_hours=1,
+        magnetic_retention_days=1,
+    )
+
+    tables_in_db = wr.timestream.list_tables(database=timestream_database_and_table)
+    assert f"{timestream_database_and_table}_2" in tables_in_db
+
+    wr.timestream.delete_table(database=timestream_database_and_table, table=f"{timestream_database_and_table}_2")
+
+    tables_in_db = wr.timestream.list_tables(database=timestream_database_and_table)
+    assert f"{timestream_database_and_table}_2" not in tables_in_db