diff --git a/llms.txt b/llms.txt index 034d50c..8f93419 100644 --- a/llms.txt +++ b/llms.txt @@ -318,6 +318,39 @@ Export data to a database. **validate_api_key()** Test API key authentication. +### Set Management + +A "set" is a persistent, named snapshot of rows (e.g. a regression dataset or a +shared fixture). Sets are scoped to a project/team and appear on the DataMaker +Sets page. + +**get_sets(project_id=None)** +Get all saved sets. Pass `project_id` (or set `DATAMAKER_PROJECT_ID`) to scope +the listing to one project. + +**get_set(set_id)** +Get a single set by ID, including its full saved rows in `data`. + +**create_set(name, data=None, description=None, row_count=None, project_id=None)** +Create (save) a new set. `data` is typically a list of row dicts. + +**update_set(set_id, name=None, description=None, data=None, row_count=None)** +Update a saved set. Only the provided fields are changed. + +**delete_set(set_id)** +Delete a saved set by ID. + +**save_set(name, data, description=None, project_id=None)** +Convenience method to save rows as a named set. Pulls project context from +`DATAMAKER_PROJECT_ID` when not provided. + +```python +dm = DataMaker() +rows = dm.generate(template) +result = dm.save_set("nightly-regression", rows) +print(f"Saved set: {result['name']} ({result['rowCount']} rows)") +``` + ## Field Types Reference ### Basic Types diff --git a/src/datamaker/main.py b/src/datamaker/main.py index 3bcdb75..ddf757b 100644 --- a/src/datamaker/main.py +++ b/src/datamaker/main.py @@ -21,6 +21,7 @@ ) from .routes.export_and_validation import ExportClient, ValidationClient from .routes.scenario_files import ScenarioFilesClient +from .routes.sets import SetsClient load_dotenv() @@ -68,6 +69,7 @@ def __init__( self._scenario_files = ScenarioFilesClient( api_key, default_headers, base_url, verify ) + self._sets = SetsClient(api_key, default_headers, base_url, verify) # Maintain backward compatibility self.api_key = self._generation.api_key @@ -661,6 +663,131 @@ def save_file( folder=folder, ) + # =================== SET METHODS =================== + def get_sets(self, project_id: Optional[str] = None): + """Get all saved sets for the caller's project/team scope. + + Args: + project_id: Optional project ID to scope the listing to. Falls back + to the DATAMAKER_PROJECT_ID env var. + + Returns: + A list of set metadata dictionaries. + """ + return self._sets.get_sets(project_id) + + def get_set(self, set_id: str): + """Get a single saved set by ID, including its full rows payload. + + Args: + set_id: The unique identifier of the set. + + Returns: + The set dictionary (including its saved rows in ``data``). + """ + return self._sets.get_set(set_id) + + def create_set( + self, + name: str, + data=None, + description: Optional[str] = None, + row_count: Optional[int] = None, + project_id: Optional[str] = None, + ): + """Create (save) a new set. + + Args: + name: A name for the saved set. + data: The rows payload to save - typically a list of row dicts. + description: Optional short description of what the set captures. + row_count: Optional explicit row count (derived from data otherwise). + project_id: Optional project ID. Falls back to DATAMAKER_PROJECT_ID. + + Returns: + The created set dictionary. + """ + return self._sets.create_set( + name=name, + data=data, + description=description, + row_count=row_count, + project_id=project_id, + ) + + def update_set( + self, + set_id: str, + name: Optional[str] = None, + description: Optional[str] = None, + data=None, + row_count: Optional[int] = None, + ): + """Update a saved set. + + Args: + set_id: The unique identifier of the set to update. + name: New name for the set. + description: New description. + data: Replacement rows payload. + row_count: Explicit row count override. + + Returns: + The updated set dictionary. + """ + return self._sets.update_set( + set_id, + name=name, + description=description, + data=data, + row_count=row_count, + ) + + def delete_set(self, set_id: str): + """Delete a saved set by ID. + + Args: + set_id: The unique identifier of the set to delete. + + Returns: + Confirmation response. + """ + return self._sets.delete_set(set_id) + + def save_set( + self, + name: str, + data, + description: Optional[str] = None, + project_id: Optional[str] = None, + ): + """Convenience method to save rows as a named set. + + Project context is pulled from DATAMAKER_PROJECT_ID when not provided, + which is convenient inside DataMaker sandbox environments. + + Args: + name: A name for the saved set. + data: The rows to save - typically a list of row dicts. + description: Optional short description of what the set captures. + project_id: Optional project ID. Falls back to DATAMAKER_PROJECT_ID. + + Returns: + The created set dictionary. + + Example: + >>> dm = DataMaker() + >>> rows = dm.generate(template) + >>> result = dm.save_set("nightly-regression", rows) + >>> print(f"Saved set: {result['name']} ({result['rowCount']} rows)") + """ + return self._sets.save_set( + name=name, + data=data, + description=description, + project_id=project_id, + ) + # =================== PROPERTY ACCESS TO CLIENTS =================== # For advanced users who want direct access to specific clients @property @@ -747,3 +874,8 @@ def validation(self): def scenario_files(self): """Access to scenario files client.""" return self._scenario_files + + @property + def sets(self): + """Access to sets client.""" + return self._sets diff --git a/src/datamaker/routes/__init__.py b/src/datamaker/routes/__init__.py index 1509291..9f6368a 100644 --- a/src/datamaker/routes/__init__.py +++ b/src/datamaker/routes/__init__.py @@ -12,6 +12,7 @@ from .folders_and_utils import TemplateFoldersClient, ShortcutsClient, FeedbackClient from .export_and_validation import ExportClient, ValidationClient from .scenario_files import ScenarioFilesClient +from .sets import SetsClient __all__ = [ "BaseClient", @@ -32,4 +33,5 @@ "ExportClient", "ValidationClient", "ScenarioFilesClient", + "SetsClient", ] diff --git a/src/datamaker/routes/sets.py b/src/datamaker/routes/sets.py new file mode 100644 index 0000000..875805f --- /dev/null +++ b/src/datamaker/routes/sets.py @@ -0,0 +1,177 @@ +"""Client for set operations - persistent named snapshots of rows. + +A "set" is a saved snapshot of generated or fetched rows (e.g. a pinned +regression dataset, a shared fixture, or an audit copy of what was masked and +exported). Sets are scoped to a project/team and show up on the DataMaker Sets +page. +""" + +import os +from typing import Dict, List, Optional, Any +from .base import BaseClient +from ..error import DataMakerError + + +class SetsClient(BaseClient): + """Client for set operations (persistent saved row snapshots).""" + + def get_sets(self, project_id: Optional[str] = None) -> List[Dict]: + """Fetch all saved sets for the caller's project/team scope. + + Args: + project_id: Optional project ID to scope the listing to. Falls back + to the DATAMAKER_PROJECT_ID env var. When omitted, the API + returns every set the API key can access. + + Returns: + A list of set metadata dictionaries. + """ + project_id = project_id or os.environ.get("DATAMAKER_PROJECT_ID") + + endpoint = "/sets" + if project_id: + endpoint += f"?projectId={project_id}" + + response = self._make_request("GET", endpoint) + return response.json() + + def get_set(self, set_id: str) -> Dict: + """Get a single saved set by ID, including its full ``data`` payload. + + Args: + set_id: The unique identifier of the set. + + Returns: + The set dictionary (including its saved rows in ``data``). + """ + response = self._make_request("GET", f"/sets/{set_id}") + return response.json() + + def create_set( + self, + name: str, + data: Optional[Any] = None, + description: Optional[str] = None, + row_count: Optional[int] = None, + project_id: Optional[str] = None, + ) -> Dict: + """Create (save) a new set. + + Args: + name: A name for the saved set. + data: The rows payload to save - typically a list of row dicts. + When omitted, an empty set is created. + description: Optional short description of what the set captures. + row_count: Optional explicit row count. Derived from ``data`` by the + API when omitted. + project_id: Optional project ID the set belongs to. Falls back to the + DATAMAKER_PROJECT_ID env var. Required when the API key is not + already scoped to a single project. + + Returns: + The created set dictionary. + + Raises: + DataMakerError: If ``name`` is not provided. + """ + if not name: + raise DataMakerError("name is required to create a set.") + + project_id = project_id or os.environ.get("DATAMAKER_PROJECT_ID") + + set_data: Dict[str, Any] = {"name": name} + if data is not None: + set_data["data"] = data + if description is not None: + set_data["description"] = description + if row_count is not None: + set_data["rowCount"] = row_count + if project_id: + set_data["projectId"] = project_id + + response = self._make_request("POST", "/sets", json=set_data) + return response.json() + + def update_set( + self, + set_id: str, + name: Optional[str] = None, + description: Optional[str] = None, + data: Optional[Any] = None, + row_count: Optional[int] = None, + ) -> Dict: + """Update a saved set. + + Only the fields that are provided are sent; the rest are left unchanged. + + Args: + set_id: The unique identifier of the set to update. + name: New name for the set. + description: New description (pass ``None`` to leave unchanged). + data: Replacement rows payload. The API keeps ``rowCount`` in sync + whenever ``data`` changes. + row_count: Explicit row count override. + + Returns: + The updated set dictionary. + """ + update_data: Dict[str, Any] = {} + if name is not None: + update_data["name"] = name + if description is not None: + update_data["description"] = description + if data is not None: + update_data["data"] = data + if row_count is not None: + update_data["rowCount"] = row_count + + response = self._make_request("PATCH", f"/sets/{set_id}", json=update_data) + return response.json() + + def delete_set(self, set_id: str) -> Dict: + """Delete a saved set by ID. + + Args: + set_id: The unique identifier of the set to delete. + + Returns: + Confirmation response. + """ + response = self._make_request("DELETE", f"/sets/{set_id}") + return response.json() + + def save_set( + self, + name: str, + data: Any, + description: Optional[str] = None, + project_id: Optional[str] = None, + ) -> Dict: + """Convenience method to save rows as a named set. + + This is the most common entry point: hand it a name and the rows you + want to keep. Project context is pulled from the DATAMAKER_PROJECT_ID + env var when not provided, which is convenient inside DataMaker sandbox + environments where it is pre-configured. + + Args: + name: A name for the saved set. + data: The rows to save - typically a list of row dicts. + description: Optional short description of what the set captures. + project_id: Optional project ID. Falls back to DATAMAKER_PROJECT_ID. + + Returns: + The created set dictionary. + + Example: + >>> dm = DataMaker() + >>> rows = dm.generate(template) + >>> result = dm.save_set("nightly-regression", rows) + >>> print(f"Saved set: {result['name']} ({result['rowCount']} rows)") + """ + return self.create_set( + name=name, + data=data, + description=description, + project_id=project_id, + ) diff --git a/tests/test_routes.py b/tests/test_routes.py index a0877b5..53ea5a2 100644 --- a/tests/test_routes.py +++ b/tests/test_routes.py @@ -1,5 +1,6 @@ """Tests for the route client classes.""" +import os import pytest from unittest.mock import Mock, patch, MagicMock from src.datamaker.routes.base import BaseClient @@ -10,6 +11,7 @@ from src.datamaker.routes.projects import ProjectsClient from src.datamaker.routes.users import UsersClient from src.datamaker.routes.teams import TeamsClient, TeamMembersClient +from src.datamaker.routes.sets import SetsClient from src.datamaker.error import DataMakerError @@ -503,3 +505,162 @@ def test_invite_team_member(self, mock_make_request, api_key): "POST", "/teamMembers/invite", json=invite_data ) assert result == {"success": True, "inviteId": "invite-1"} + + +class TestSetsClient: + """Test cases for the SetsClient class.""" + + @patch("src.datamaker.routes.base.BaseClient._make_request") + def test_get_sets(self, mock_make_request, api_key): + """Test getting all sets (unscoped when no project id).""" + mock_response = Mock() + mock_response.json.return_value = [{"id": "1", "name": "Set 1"}] + mock_make_request.return_value = mock_response + + # Ensure no env var bleeds into the scope for this assertion. + with patch.dict("os.environ", {}, clear=False): + os.environ.pop("DATAMAKER_PROJECT_ID", None) + client = SetsClient(api_key=api_key) + result = client.get_sets() + + mock_make_request.assert_called_once_with("GET", "/sets") + assert result == [{"id": "1", "name": "Set 1"}] + + @patch("src.datamaker.routes.base.BaseClient._make_request") + def test_get_sets_with_project_id(self, mock_make_request, api_key): + """Test getting sets scoped to a project id.""" + mock_response = Mock() + mock_response.json.return_value = [{"id": "1", "name": "Set 1"}] + mock_make_request.return_value = mock_response + + client = SetsClient(api_key=api_key) + result = client.get_sets(project_id="proj-1") + + mock_make_request.assert_called_once_with("GET", "/sets?projectId=proj-1") + assert result == [{"id": "1", "name": "Set 1"}] + + @patch("src.datamaker.routes.base.BaseClient._make_request") + def test_get_set(self, mock_make_request, api_key): + """Test getting a single set by id.""" + mock_response = Mock() + mock_response.json.return_value = {"id": "1", "name": "Set 1", "data": []} + mock_make_request.return_value = mock_response + + client = SetsClient(api_key=api_key) + result = client.get_set("1") + + mock_make_request.assert_called_once_with("GET", "/sets/1") + assert result == {"id": "1", "name": "Set 1", "data": []} + + @patch("src.datamaker.routes.base.BaseClient._make_request") + def test_create_set(self, mock_make_request, api_key): + """Test creating a set with rows + project id.""" + mock_response = Mock() + mock_response.json.return_value = {"id": "1", "name": "New Set", "rowCount": 2} + mock_make_request.return_value = mock_response + + client = SetsClient(api_key=api_key) + rows = [{"a": 1}, {"a": 2}] + result = client.create_set( + name="New Set", + data=rows, + description="A test set", + project_id="proj-1", + ) + + expected_data = { + "name": "New Set", + "data": rows, + "description": "A test set", + "projectId": "proj-1", + } + mock_make_request.assert_called_once_with("POST", "/sets", json=expected_data) + assert result == {"id": "1", "name": "New Set", "rowCount": 2} + + @patch("src.datamaker.routes.base.BaseClient._make_request") + def test_create_set_minimal(self, mock_make_request, api_key): + """Test creating a set with only a name (no optional fields).""" + mock_response = Mock() + mock_response.json.return_value = {"id": "1", "name": "Empty Set"} + mock_make_request.return_value = mock_response + + with patch.dict("os.environ", {}, clear=False): + os.environ.pop("DATAMAKER_PROJECT_ID", None) + client = SetsClient(api_key=api_key) + result = client.create_set(name="Empty Set") + + mock_make_request.assert_called_once_with( + "POST", "/sets", json={"name": "Empty Set"} + ) + assert result == {"id": "1", "name": "Empty Set"} + + def test_create_set_requires_name(self, api_key): + """Test that creating a set without a name raises an error.""" + client = SetsClient(api_key=api_key) + with pytest.raises(DataMakerError, match="name is required"): + client.create_set(name="") + + @patch("src.datamaker.routes.base.BaseClient._make_request") + def test_create_set_uses_env_project_id(self, mock_make_request, api_key): + """Test that create_set falls back to DATAMAKER_PROJECT_ID.""" + mock_response = Mock() + mock_response.json.return_value = {"id": "1", "name": "Env Set"} + mock_make_request.return_value = mock_response + + with patch.dict("os.environ", {"DATAMAKER_PROJECT_ID": "env-proj"}): + client = SetsClient(api_key=api_key) + client.create_set(name="Env Set", data=[{"x": 1}]) + + expected_data = {"name": "Env Set", "data": [{"x": 1}], "projectId": "env-proj"} + mock_make_request.assert_called_once_with("POST", "/sets", json=expected_data) + + @patch("src.datamaker.routes.base.BaseClient._make_request") + def test_update_set(self, mock_make_request, api_key): + """Test updating a set sends only the provided fields.""" + mock_response = Mock() + mock_response.json.return_value = {"id": "1", "name": "Renamed"} + mock_make_request.return_value = mock_response + + client = SetsClient(api_key=api_key) + result = client.update_set("1", name="Renamed", data=[{"a": 1}]) + + expected_data = {"name": "Renamed", "data": [{"a": 1}]} + mock_make_request.assert_called_once_with( + "PATCH", "/sets/1", json=expected_data + ) + assert result == {"id": "1", "name": "Renamed"} + + @patch("src.datamaker.routes.base.BaseClient._make_request") + def test_delete_set(self, mock_make_request, api_key): + """Test deleting a set by id.""" + mock_response = Mock() + mock_response.json.return_value = {"message": "Set deleted"} + mock_make_request.return_value = mock_response + + client = SetsClient(api_key=api_key) + result = client.delete_set("1") + + mock_make_request.assert_called_once_with("DELETE", "/sets/1") + assert result == {"message": "Set deleted"} + + @patch("src.datamaker.routes.base.BaseClient._make_request") + def test_save_set(self, mock_make_request, api_key): + """Test the save_set convenience method delegates to create_set.""" + mock_response = Mock() + mock_response.json.return_value = { + "id": "1", + "name": "Snapshot", + "rowCount": 1, + } + mock_make_request.return_value = mock_response + + client = SetsClient(api_key=api_key) + result = client.save_set("Snapshot", [{"a": 1}], project_id="proj-1") + + expected_data = { + "name": "Snapshot", + "data": [{"a": 1}], + "projectId": "proj-1", + } + mock_make_request.assert_called_once_with("POST", "/sets", json=expected_data) + assert result["rowCount"] == 1