diff --git a/src/datamaker/main.py b/src/datamaker/main.py index e5d7009..c07e8a7 100644 --- a/src/datamaker/main.py +++ b/src/datamaker/main.py @@ -1,6 +1,6 @@ import os from dotenv import load_dotenv -from typing import Optional, Dict +from typing import Optional, Dict, List from .routes.base import BaseClient from .routes.generation import GenerationClient from .routes.templates import TemplatesClient @@ -22,6 +22,7 @@ from .routes.export_and_validation import ExportClient, ValidationClient from .routes.scenario_files import ScenarioFilesClient from .routes.sets import SetsClient +from .routes.keymaps import KeyMapsClient load_dotenv() @@ -70,6 +71,7 @@ def __init__( api_key, default_headers, base_url, verify ) self._sets = SetsClient(api_key, default_headers, base_url, verify) + self._keymaps = KeyMapsClient(api_key, default_headers, base_url, verify) # Maintain backward compatibility self.api_key = self._generation.api_key @@ -793,6 +795,141 @@ def save_set( project_id=project_id, ) + def get_keymaps(self, project_id: Optional[str] = None): + """List key maps for a project: one row per (mapName, object). + + Args: + project_id: Optional project ID. Falls back to DATAMAKER_PROJECT_ID. + + Returns: + A list of dictionaries with mapName, object, entryCount, updatedAt. + """ + return self._keymaps.get_keymaps(project_id) + + def keymap_put( + self, + map_name: str, + object: str, + entries: Dict[str, str], + run_id: Optional[str] = None, + project_id: Optional[str] = None, + ): + """Record old-to-new key mappings in a named key map (batch upsert). + + Use after creating records in a target system to remember which source + key became which target key. Writing the same old key again overwrites + its new key. Batches are capped at 5000 entries per call. + + Args: + map_name: Logical map name, e.g. "sap-material-migration". + object: The domain object type, e.g. "Material". + entries: Mapping of old key to new key. + run_id: Optional run/job id that minted these keys. + project_id: Optional project ID. Falls back to DATAMAKER_PROJECT_ID. + + Returns: + A dictionary with mapName, object and upserted (count). + + Example: + >>> dm = DataMaker() + >>> dm.keymap_put( + ... "sap-material-migration", + ... "Material", + ... {"MAT-001": "700001", "MAT-002": "700002"}, + ... ) + """ + return self._keymaps.keymap_put( + map_name=map_name, + object=object, + entries=entries, + run_id=run_id, + project_id=project_id, + ) + + def keymap_lookup( + self, + map_name: str, + object: str, + old_keys: List[str], + project_id: Optional[str] = None, + ): + """Translate source-system keys to target-system keys (batch lookup). + + Use when generating or migrating data that references records migrated + earlier. Lookups are capped at 5000 keys per call. + + Args: + map_name: The key map to look up in. + object: The domain object type, e.g. "Material". + old_keys: Source-system keys to translate. + project_id: Optional project ID. Falls back to DATAMAKER_PROJECT_ID. + + Returns: + A dictionary with mappings (found) and missing (no mapping yet). + + Example: + >>> dm = DataMaker() + >>> result = dm.keymap_lookup( + ... "sap-material-migration", "Material", ["MAT-001", "MAT-999"] + ... ) + """ + return self._keymaps.keymap_lookup( + map_name=map_name, + object=object, + old_keys=old_keys, + project_id=project_id, + ) + + def get_keymap_entries( + self, + map_name: str, + object: Optional[str] = None, + page: int = 1, + page_size: int = 100, + project_id: Optional[str] = None, + ): + """Fetch a page of a key map's entries for inspection. + + Args: + map_name: The key map to read. + object: Optional domain object type filter. + page: 1-based page number. + page_size: Entries per page (server-capped at 500). + project_id: Optional project ID. Falls back to DATAMAKER_PROJECT_ID. + + Returns: + A dictionary with entries, total, page and pageSize. + """ + return self._keymaps.get_keymap_entries( + map_name=map_name, + object=object, + page=page, + page_size=page_size, + project_id=project_id, + ) + + def delete_keymap( + self, + map_name: str, + object: Optional[str] = None, + project_id: Optional[str] = None, + ): + """Drop a key map (all its entries), optionally one object type only. + + Args: + map_name: The key map to delete. + object: Optional domain object type to scope the delete to. + project_id: Optional project ID. Falls back to DATAMAKER_PROJECT_ID. + + Returns: + Confirmation response with the deleted entry count. + """ + return self._keymaps.delete_keymap( + map_name=map_name, + object=object, + project_id=project_id, + ) + # =================== PROPERTY ACCESS TO CLIENTS =================== # For advanced users who want direct access to specific clients @property @@ -884,3 +1021,8 @@ def scenario_files(self): def sets(self): """Access to sets client.""" return self._sets + + @property + def keymaps(self): + """Access to key maps client.""" + return self._keymaps diff --git a/src/datamaker/routes/__init__.py b/src/datamaker/routes/__init__.py index 9f6368a..9c7eca8 100644 --- a/src/datamaker/routes/__init__.py +++ b/src/datamaker/routes/__init__.py @@ -13,6 +13,7 @@ from .export_and_validation import ExportClient, ValidationClient from .scenario_files import ScenarioFilesClient from .sets import SetsClient +from .keymaps import KeyMapsClient __all__ = [ "BaseClient", @@ -34,4 +35,5 @@ "ValidationClient", "ScenarioFilesClient", "SetsClient", + "KeyMapsClient", ] diff --git a/src/datamaker/routes/keymaps.py b/src/datamaker/routes/keymaps.py new file mode 100644 index 0000000..e51a8b0 --- /dev/null +++ b/src/datamaker/routes/keymaps.py @@ -0,0 +1,221 @@ +"""Client for key map operations - old-to-new key mappings for migrations. + +A "key map" records which source-system key became which target-system key +during a migration (e.g. legacy material number to new material number). +Unlike a set (one JSON blob), every mapping is its own row on the server, so +lookups are indexed and batch writes from parallel load workers are +concurrency-safe. Entries are unique per (project, map name, object, old key) +and the last write wins for the new key. +""" + +import os +from typing import Dict, List, Optional +from .base import BaseClient +from ..error import DataMakerError + + +class KeyMapsClient(BaseClient): + """Client for key map operations (old-to-new key mappings).""" + + def get_keymaps(self, project_id: Optional[str] = None) -> List[Dict]: + """List key maps for a project: one row per (mapName, object). + + Args: + project_id: Optional project ID to scope the listing to. Falls back + to the DATAMAKER_PROJECT_ID env var. + + Returns: + A list of dictionaries with ``mapName``, ``object``, ``entryCount`` + and ``updatedAt``. + """ + project_id = project_id or os.environ.get("DATAMAKER_PROJECT_ID") + + endpoint = "/keymaps" + if project_id: + endpoint += f"?projectId={project_id}" + + response = self._make_request("GET", endpoint) + return response.json() + + def keymap_put( + self, + map_name: str, + object: str, + entries: Dict[str, str], + run_id: Optional[str] = None, + project_id: Optional[str] = None, + ) -> Dict: + """Record old-to-new key mappings in a named key map (batch upsert). + + Use after creating records in a target system to remember which source + key became which target key, so dependent data can reference the right + keys later. Writing the same old key again overwrites its new key. + Batches are capped at 5000 entries per call - page larger writes. + + Args: + map_name: Logical map name grouping the entries, e.g. + "sap-material-migration". + object: The domain object type the keys belong to, e.g. "Material" + or "BusinessPartner". + entries: Mapping of old key to new key. + run_id: Optional run/job id that minted these keys. + project_id: Optional project ID. Falls back to DATAMAKER_PROJECT_ID. + + Returns: + A dictionary with ``mapName``, ``object`` and ``upserted`` (count). + + Raises: + DataMakerError: If ``map_name``, ``object`` or ``entries`` is empty. + + Example: + >>> dm = DataMaker() + >>> dm.keymap_put( + ... "sap-material-migration", + ... "Material", + ... {"MAT-001": "700001", "MAT-002": "700002"}, + ... ) + """ + if not map_name: + raise DataMakerError("map_name is required to put key map entries.") + if not object: + raise DataMakerError("object is required to put key map entries.") + if not entries: + raise DataMakerError("entries must not be empty.") + + project_id = project_id or os.environ.get("DATAMAKER_PROJECT_ID") + + payload: Dict = { + "mapName": map_name, + "object": object, + "entries": [ + {"oldKey": old_key, "newKey": new_key} + for old_key, new_key in entries.items() + ], + } + if run_id is not None: + payload["runId"] = run_id + if project_id: + payload["projectId"] = project_id + + response = self._make_request("POST", "/keymaps/entries", json=payload) + return response.json() + + def keymap_lookup( + self, + map_name: str, + object: str, + old_keys: List[str], + project_id: Optional[str] = None, + ) -> Dict: + """Translate source-system keys to target-system keys (batch lookup). + + Use when generating or migrating data that references records migrated + earlier (e.g. orders that need the NEW material numbers for OLD ones). + Lookups are capped at 5000 keys per call - page larger reads. + + Args: + map_name: The key map to look up in. + object: The domain object type, e.g. "Material". + old_keys: Source-system keys to translate. + project_id: Optional project ID. Falls back to DATAMAKER_PROJECT_ID. + + Returns: + A dictionary with ``mappings`` (old key to new key for the keys + that were found) and ``missing`` (keys with no mapping yet). + + Raises: + DataMakerError: If ``map_name``, ``object`` or ``old_keys`` is empty. + + Example: + >>> dm = DataMaker() + >>> result = dm.keymap_lookup( + ... "sap-material-migration", "Material", ["MAT-001", "MAT-999"] + ... ) + >>> result["mappings"] + {'MAT-001': '700001'} + >>> result["missing"] + ['MAT-999'] + """ + if not map_name: + raise DataMakerError("map_name is required to look up key mappings.") + if not object: + raise DataMakerError("object is required to look up key mappings.") + if not old_keys: + raise DataMakerError("old_keys must not be empty.") + + project_id = project_id or os.environ.get("DATAMAKER_PROJECT_ID") + + payload: Dict = { + "mapName": map_name, + "object": object, + "oldKeys": old_keys, + } + if project_id: + payload["projectId"] = project_id + + response = self._make_request("POST", "/keymaps/lookup", json=payload) + return response.json() + + def get_keymap_entries( + self, + map_name: str, + object: Optional[str] = None, + page: int = 1, + page_size: int = 100, + project_id: Optional[str] = None, + ) -> Dict: + """Fetch a page of a key map's entries for inspection. + + Args: + map_name: The key map to read. + object: Optional domain object type filter. + page: 1-based page number. + page_size: Entries per page (server-capped at 500). + project_id: Optional project ID. Falls back to DATAMAKER_PROJECT_ID. + + Returns: + A dictionary with ``entries``, ``total``, ``page`` and ``pageSize``. + """ + project_id = project_id or os.environ.get("DATAMAKER_PROJECT_ID") + + params = [f"page={page}", f"pageSize={page_size}"] + if object: + params.append(f"object={object}") + if project_id: + params.append(f"projectId={project_id}") + + endpoint = f"/keymaps/{map_name}/entries?" + "&".join(params) + response = self._make_request("GET", endpoint) + return response.json() + + def delete_keymap( + self, + map_name: str, + object: Optional[str] = None, + project_id: Optional[str] = None, + ) -> Dict: + """Drop a key map (all its entries). + + Args: + map_name: The key map to delete. + object: Optional domain object type - when given, only that + object's entries are dropped. + project_id: Optional project ID. Falls back to DATAMAKER_PROJECT_ID. + + Returns: + Confirmation response with the deleted entry count. + """ + project_id = project_id or os.environ.get("DATAMAKER_PROJECT_ID") + + params = [] + if object: + params.append(f"object={object}") + if project_id: + params.append(f"projectId={project_id}") + + endpoint = f"/keymaps/{map_name}" + if params: + endpoint += "?" + "&".join(params) + + response = self._make_request("DELETE", endpoint) + return response.json() diff --git a/tests/test_routes.py b/tests/test_routes.py index 0b512e6..69a817e 100644 --- a/tests/test_routes.py +++ b/tests/test_routes.py @@ -12,6 +12,7 @@ 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.routes.keymaps import KeyMapsClient from src.datamaker.routes.custom_types import EndpointsClient from src.datamaker.error import DataMakerError @@ -690,3 +691,134 @@ def test_resolve_endpoint_auth(self, mock_make_request, api_key): ) assert result["authType"] == "Basic" assert result["basic"]["username"] == "user" + +class TestKeyMapsClient: + """Test cases for the KeyMapsClient class.""" + + @patch("src.datamaker.routes.base.BaseClient._make_request") + def test_get_keymaps_with_project_id(self, mock_make_request, api_key): + """Test listing key maps scoped to a project id.""" + mock_response = Mock() + mock_response.json.return_value = [ + {"mapName": "m1", "object": "Material", "entryCount": 2} + ] + mock_make_request.return_value = mock_response + + client = KeyMapsClient(api_key=api_key) + result = client.get_keymaps(project_id="proj-1") + + mock_make_request.assert_called_once_with("GET", "/keymaps?projectId=proj-1") + assert result == [{"mapName": "m1", "object": "Material", "entryCount": 2}] + + @patch("src.datamaker.routes.base.BaseClient._make_request") + def test_keymap_put(self, mock_make_request, api_key): + """Test batch-upserting key map entries.""" + mock_response = Mock() + mock_response.json.return_value = { + "mapName": "sap-material-migration", + "object": "Material", + "upserted": 2, + } + mock_make_request.return_value = mock_response + + client = KeyMapsClient(api_key=api_key) + result = client.keymap_put( + "sap-material-migration", + "Material", + {"MAT-001": "700001", "MAT-002": "700002"}, + run_id="run-1", + project_id="proj-1", + ) + + mock_make_request.assert_called_once_with( + "POST", + "/keymaps/entries", + json={ + "mapName": "sap-material-migration", + "object": "Material", + "entries": [ + {"oldKey": "MAT-001", "newKey": "700001"}, + {"oldKey": "MAT-002", "newKey": "700002"}, + ], + "runId": "run-1", + "projectId": "proj-1", + }, + ) + assert result["upserted"] == 2 + + @patch("src.datamaker.routes.base.BaseClient._make_request") + def test_keymap_put_requires_entries(self, mock_make_request, api_key): + """Test that an empty entries dict raises.""" + client = KeyMapsClient(api_key=api_key) + with pytest.raises(DataMakerError): + client.keymap_put("m1", "Material", {}) + mock_make_request.assert_not_called() + + @patch("src.datamaker.routes.base.BaseClient._make_request") + def test_keymap_lookup(self, mock_make_request, api_key): + """Test batch key lookup returns mappings and missing keys.""" + mock_response = Mock() + mock_response.json.return_value = { + "mapName": "m1", + "object": "Material", + "mappings": {"MAT-001": "700001"}, + "missing": ["MAT-999"], + } + mock_make_request.return_value = mock_response + + client = KeyMapsClient(api_key=api_key) + result = client.keymap_lookup( + "m1", "Material", ["MAT-001", "MAT-999"], project_id="proj-1" + ) + + mock_make_request.assert_called_once_with( + "POST", + "/keymaps/lookup", + json={ + "mapName": "m1", + "object": "Material", + "oldKeys": ["MAT-001", "MAT-999"], + "projectId": "proj-1", + }, + ) + assert result["mappings"] == {"MAT-001": "700001"} + assert result["missing"] == ["MAT-999"] + + @patch("src.datamaker.routes.base.BaseClient._make_request") + def test_get_keymap_entries(self, mock_make_request, api_key): + """Test paginated entry inspection with an object filter.""" + mock_response = Mock() + mock_response.json.return_value = { + "mapName": "m1", + "page": 1, + "pageSize": 50, + "total": 1, + "entries": [{"object": "Material", "oldKey": "a", "newKey": "b"}], + } + mock_make_request.return_value = mock_response + + client = KeyMapsClient(api_key=api_key) + result = client.get_keymap_entries( + "m1", object="Material", page=1, page_size=50, project_id="proj-1" + ) + + mock_make_request.assert_called_once_with( + "GET", + "/keymaps/m1/entries?page=1&pageSize=50&object=Material&projectId=proj-1", + ) + assert result["total"] == 1 + + @patch("src.datamaker.routes.base.BaseClient._make_request") + def test_delete_keymap(self, mock_make_request, api_key): + """Test dropping a key map scoped to one object type.""" + mock_response = Mock() + mock_response.json.return_value = {"message": "Key map deleted", "deleted": 5} + mock_make_request.return_value = mock_response + + client = KeyMapsClient(api_key=api_key) + result = client.delete_keymap("m1", object="Material", project_id="proj-1") + + mock_make_request.assert_called_once_with( + "DELETE", "/keymaps/m1?object=Material&projectId=proj-1" + ) + assert result["deleted"] == 5