From 02a40e814d84a884afd0687e9e863971eb44fab7 Mon Sep 17 00:00:00 2001 From: Santhosh Ramaraj Date: Fri, 25 Oct 2024 11:10:36 -0500 Subject: [PATCH 1/4] feat: Generator variant of query_tables --- .../clients/dataframe/_data_frame_client.py | 29 ++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/nisystemlink/clients/dataframe/_data_frame_client.py b/nisystemlink/clients/dataframe/_data_frame_client.py index 180d8708..9efcab86 100644 --- a/nisystemlink/clients/dataframe/_data_frame_client.py +++ b/nisystemlink/clients/dataframe/_data_frame_client.py @@ -1,6 +1,6 @@ """Implementation of DataFrameClient.""" -from typing import List, Optional +from typing import Generator, List, Optional from nisystemlink.clients import core from nisystemlink.clients.core._uplink._base_client import BaseClient @@ -15,6 +15,8 @@ from requests.models import Response from uplink import Body, Field, Path, Query +from nisystemlink.clients.dataframe.models._table_metadata import TableMetadata + from . import models @@ -119,6 +121,31 @@ def query_tables(self, query: models.QueryTablesRequest) -> models.PagedTables: """ ... + def query_tables_generator( + self, query: models.QueryTablesRequest + ) -> Generator[models.TableMetadata, None, None]: + """Queries available tables on the SystemLink DataFrame service and returns their metadata. + + Args: + query: The request to query tables. `continuation_token` is ignored. + + Yields: + `models.TableMetadata` Table Metadata + """ + _query = query.copy() + _query.continuation_token = None + while True: + paged_tables = self.query_tables(query=_query) + cont_token = paged_tables.continuation_token + + for table in paged_tables.tables: + yield table + + if cont_token is None: + break + else: + _query.continuation_token = cont_token + @get("tables/{id}") def get_table_metadata(self, id: str) -> models.TableMetadata: """Retrieves the metadata and column information for a single table identified by its ID. From f9b97d9169f5d716adad2373373a6b60517c72b0 Mon Sep 17 00:00:00 2001 From: Santhosh Ramaraj Date: Fri, 25 Oct 2024 12:07:38 -0500 Subject: [PATCH 2/4] refactor: lint issues --- nisystemlink/clients/dataframe/_data_frame_client.py | 1 - 1 file changed, 1 deletion(-) diff --git a/nisystemlink/clients/dataframe/_data_frame_client.py b/nisystemlink/clients/dataframe/_data_frame_client.py index 9efcab86..d91f34ed 100644 --- a/nisystemlink/clients/dataframe/_data_frame_client.py +++ b/nisystemlink/clients/dataframe/_data_frame_client.py @@ -15,7 +15,6 @@ from requests.models import Response from uplink import Body, Field, Path, Query -from nisystemlink.clients.dataframe.models._table_metadata import TableMetadata from . import models From c0802122971ba686c0e3b04cb2d238c98247b527 Mon Sep 17 00:00:00 2001 From: Santhosh Ramaraj Date: Fri, 25 Oct 2024 16:07:22 -0500 Subject: [PATCH 3/4] tests: test for query_tables_generator --- tests/integration/dataframe/test_dataframe.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/integration/dataframe/test_dataframe.py b/tests/integration/dataframe/test_dataframe.py index aebe6147..986ab82b 100644 --- a/tests/integration/dataframe/test_dataframe.py +++ b/tests/integration/dataframe/test_dataframe.py @@ -175,6 +175,25 @@ def test__query_tables__returns( assert len(second_page.tables) == 1 assert second_page.continuation_token is None + def test__query_tables_generator__returns( + self, client: DataFrameClient, test_tables: List[str] + ): + query = QueryTablesRequest( + filter="""(id == @0 or id == @1 or id == @2) + and createdWithin <= RelativeTime.CurrentWeek + and supportsAppend == @3 and rowCount < @4""", + substitutions=[test_tables[0], test_tables[1], test_tables[2], True, 1], + reference_time=datetime.now(tz=timezone.utc), + take=2, + order_by="NAME", + order_by_descending=True, + ) + # Make the generator yield all values + tables = list(client.query_tables_generator(query)) + + assert len(tables) == 3 + assert tables[0].id == test_tables[-1] # Asserts descending order + def test__modify_table__returns(self, client: DataFrameClient, create_table): id = create_table(basic_table_model) From 4a8ca97c0325f81ac4bab9e20a0fd7d5dee966fe Mon Sep 17 00:00:00 2001 From: Santhosh Ramaraj Date: Fri, 25 Oct 2024 16:46:45 -0500 Subject: [PATCH 4/4] feat: example to use query_table_generator --- examples/dataframe/query_tables.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) create mode 100644 examples/dataframe/query_tables.py diff --git a/examples/dataframe/query_tables.py b/examples/dataframe/query_tables.py new file mode 100644 index 00000000..5f6dc75f --- /dev/null +++ b/examples/dataframe/query_tables.py @@ -0,0 +1,13 @@ +from nisystemlink.clients.dataframe import DataFrameClient +from nisystemlink.clients.dataframe.models import QueryTablesRequest + +client = DataFrameClient() + +# List all tables with more than 100k rows +query = QueryTablesRequest( + filter="rowCount > 100000", order_by="CREATED_AT", order_by_descending=True +) + +# Print the query results +for table in client.query_tables_generator(query): + print(f"{table.created_at} \t {table.name} \t Rows={table.row_count}")