Add cli command to list datasets (#517)

* Add command line to list datasets

* Add option to report only datasets for a specific DID Finder

* Increase unit test coverage
- Add a file for testing servicex_client. It's only at 92%, but we can add more tests
- Test the case of auth error in get_datasets

Author: BenGalewsky

docs/command_line.rst

@@ -76,3 +76,14 @@ clear
 Clear all of the transforms from the cache. Add ``-y`` to force the
 operation without confirming with the console.
 
+datasets
+~~~~~~~~
+
+These commands interact with datasets cached on the server
+
+list
+^^^^
+List all of the datasets cached on the server. Accepts a command line argument
+of ``--did-finder`` to filter the list of datasets by a specific DID finder such
+as ``rucio`` or ``user``.
+

servicex/app/datasets.py

@@ -0,0 +1,79 @@
+# Copyright (c) 2024, IRIS-HEP
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+import asyncio
+from typing import Optional
+
+import rich
+
+from servicex.app.cli_options import url_cli_option, backend_cli_option
+
+import typer
+
+from servicex.servicex_client import ServiceXClient
+from rich.table import Table
+
+datasets_app = typer.Typer(name="datasets", no_args_is_help=True)
+
+
+@datasets_app.command(no_args_is_help=True)
+def list(
+        url: Optional[str] = url_cli_option,
+        backend: Optional[str] = backend_cli_option,
+        did_finder: Optional[str] = typer.Option(
+            None,
+            help="Filter datasets by DID finder. Some useful values are 'rucio' or 'user'",
+            show_default=False,
+        ),
+):
+    """
+    List the datasets.
+    """
+    sx = ServiceXClient(url=url, backend=backend)
+    table = Table(title="ServiceX Datasets")
+    table.add_column("ID")
+    table.add_column("Name")
+    table.add_column("Files")
+    table.add_column("Size")
+    table.add_column("Status")
+    table.add_column("Created")
+    datasets = asyncio.run(sx.get_datasets(did_finder=did_finder))
+    for d in datasets:
+        # Format the CachedDataset object into a table row
+        # The last_updated field is what we should be displaying, but that is
+        # currently set to 1970-00-00 in the server and is never updated.
+        # Stick with the last_used field until
+        # https://github.com/ssl-hep/ServiceX/issues/906 is resolved
+        table.add_row(
+            str(d.id),
+            d.name if d.did_finder != "user" else "File list",
+            "%d" % d.n_files,
+            "{:,}MB".format(round(d.size / 1e6)),
+            d.lookup_status,
+            d.last_used.strftime('%Y-%m-%dT%H:%M:%S'),
+        )
+    rich.print(table)

servicex/app/main.py

@@ -31,6 +31,7 @@
 import typer
 
 from servicex._version import __version__
+from servicex.app.datasets import datasets_app
 from servicex.app.transforms import transforms_app
 from servicex.app.cache import cache_app
 from servicex.app.codegen import codegen_app
@@ -40,6 +41,7 @@
 app.add_typer(transforms_app)
 app.add_typer(cache_app)
 app.add_typer(codegen_app)
+app.add_typer(datasets_app)
 
 
 def show_version(show: bool):

servicex/models.py

@@ -170,3 +170,18 @@ class TransformedResults(BaseModel):
     files: int
     result_format: ResultFormat
     log_url: Optional[str] = None
+
+
+class CachedDataset(BaseModel):
+    """
+    Model for a cached dataset held by ServiceX server
+    """
+    id: int
+    name: str
+    did_finder: str
+    n_files: int
+    size: int
+    events: int
+    last_used: datetime
+    last_updated: datetime
+    lookup_status: str

servicex/servicex_adapter.py

@@ -34,7 +34,7 @@
 from google.auth import jwt
 from tenacity import AsyncRetrying, stop_after_attempt, wait_fixed, retry_if_not_exception_type
 
-from servicex.models import TransformRequest, TransformStatus
+from servicex.models import TransformRequest, TransformStatus, CachedDataset
 
 
 class AuthorizationError(BaseException):
@@ -109,6 +109,22 @@ def get_code_generators(self):
                     f"Not authorized to access serviceX at {self.url}")
             return r.json()
 
+    async def get_datasets(self, did_finder=None) -> List[CachedDataset]:
+        headers = await self._get_authorization()
+
+        with httpx.Client() as client:
+            params = {"did-finder": did_finder} if did_finder else {}
+            r = client.get(headers=headers,
+                           url=f"{self.url}/servicex/datasets",
+                           params=params)
+
+            if r.status_code == 403:
+                raise AuthorizationError(
+                    f"Not authorized to access serviceX at {self.url}")
+
+            datasets = [CachedDataset(**d) for d in r.json()['datasets']]
+            return datasets
+
     async def submit_transform(self, transform_request: TransformRequest):
         headers = await self._get_authorization()
         retry_options = ExponentialRetry(attempts=3, start_timeout=30)

servicex/servicex_client.py

@@ -280,6 +280,13 @@ async def get_transform_status_async(self, transform_id) -> TransformStatus:
 
     get_transform_status = make_sync(get_transform_status_async)
 
+    def get_datasets(self, did_finder=None):
+        r"""
+        Retrieve all datasets you have run on the server
+        :return: List of Query objects
+        """
+        return self.servicex.get_datasets(did_finder)
+
     def get_code_generators(self, backend=None):
         r"""
         Retrieve the code generators deployed with the serviceX instance

tests/test_servicex_adapter.py

@@ -129,6 +129,39 @@ def test_get_codegens_error(get, servicex):
         assert "Not authorized to access serviceX at" in str(err.value)
 
 
+@pytest.mark.asyncio
+@patch('servicex.servicex_adapter.httpx.Client.get')
+async def test_get_datasets(get, servicex):
+    get.return_value = httpx.Response(200, json={
+        "datasets": [
+            {
+                "id": "123",
+                "name": "dataset1",
+                "events": 100,
+                "size": 1000,
+                "n_files": 1,
+                "last_used": "2022-01-01T00:00:00.000000Z",
+                "last_updated": "2022-01-01T00:00:00.000000Z",
+                "lookup_status": "looking",
+                "did_finder": "rucio"
+            }
+
+        ]
+    })
+    c = await servicex.get_datasets()
+    assert len(c) == 1
+    assert c[0].id == 123
+
+
+@pytest.mark.asyncio
+@patch('servicex.servicex_adapter.httpx.Client.get')
+async def test_get_datasets_auth_error(get, servicex):
+    get.return_value = httpx.Response(403)
+    with pytest.raises(AuthorizationError) as err:
+        await servicex.get_datasets()
+        assert "Not authorized to access serviceX at" in str(err.value)
+
+
 @pytest.mark.asyncio
 @patch('servicex.servicex_adapter.RetryClient.post')
 async def test_submit(post, servicex):

tests/test_servicex_client.py

@@ -0,0 +1,63 @@
+# Copyright (c) 2024, IRIS-HEP
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+from unittest.mock import MagicMock
+
+from pytest_asyncio import fixture
+
+from servicex.query_cache import QueryCache
+from servicex.servicex_adapter import ServiceXAdapter
+from servicex.servicex_client import ServiceXClient
+
+
+@fixture
+def servicex_adaptor(mocker):
+    adapter_mock = mocker.patch('servicex.servicex_client.ServiceXAdapter')
+    mock_adapter = MagicMock(spec=ServiceXAdapter)
+
+    adapter_mock.return_value = mock_adapter
+    return mock_adapter
+
+
+@fixture
+def mock_cache(mocker):
+    cache_mock = mocker.patch('servicex.servicex_client.QueryCache')
+    mock_cache = MagicMock(spec=QueryCache)
+    mock_cache.get_codegen_by_backend.return_value = {
+        "codegens": {
+            "ROOT": "my_root_generator",
+            "UPROOT": "my_uproot_generator"
+        }
+    }
+    cache_mock.return_value = mock_cache
+    return cache_mock
+
+
+def test_get_datasets(mock_cache, servicex_adaptor):
+    sx = ServiceXClient(config_path="tests/example_config.yaml")
+    sx.get_datasets()
+    servicex_adaptor.get_datasets.assert_called_once()