Skip to content

Commit b610a34

Browse files
committed
feat(client): add RG config CRUD and versioned generate methods (DM-3958)
- add RGConfig pydantic model and RGConfigClient mixin: list/get/create/update/delete + defaults yaml - add generate_ruleset_with_rg_config (v3) and file/async/csv variants (v2), mirroring the discovery from_config split - old generate request models reject rg_config and point callers at the new methods - document dependent-RGC behaviour on ruleset library delete/usage APIs
1 parent d7010e6 commit b610a34

10 files changed

Lines changed: 1173 additions & 25 deletions

File tree

HISTORY.rst

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,14 @@ History
1111
``DiscoveryConfigLibrary`` and the ``config_type`` argument from
1212
``get_discovery_config_library_by_name``, ``create_or_update_discovery_config_library``, and
1313
``delete_discovery_config_library_by_name_if_exists``. Requires a 3.26.14 or later server.
14+
* Added ruleset-generation (RG) config management APIs
15+
(``list_rg_configs``, ``create_rg_config``, ``get_default_rg_config_yaml``, and friends).
16+
* Added versioned ruleset-generation methods that take a required ``rg_config``
17+
(``generate_ruleset_with_rg_config``, ``generate_file_ruleset_with_rg_config``,
18+
``start_async_ruleset_generation_with_rg_config``,
19+
``start_async_ruleset_generation_from_csv_with_rg_config``).
20+
The existing generate methods keep using the prior API versions,
21+
which apply the server's default RG config.
1422

1523
1.1.8 (2026-07-27)
1624
------------------

datamasque/client/__init__.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,11 +67,13 @@
6767
FileFilter,
6868
FileFilterMatchAgainst,
6969
FileRulesetGenerationRequest,
70+
FileRulesetGenerationWithRGConfigRequest,
7071
ForeignKeyRef,
7172
InDataDiscoveryConfig,
7273
InDataDiscoveryRule,
7374
ReferencingForeignKey,
7475
RulesetGenerationRequest,
76+
RulesetGenerationWithRGConfigRequest,
7577
SchemaDiscoveryColumn,
7678
SchemaDiscoveryFromConfigRequest,
7779
SchemaDiscoveryPage,
@@ -105,6 +107,7 @@
105107
RulesetPlanUpdateRequest,
106108
)
107109
from datamasque.client.models.license import LicenseInfo, SwitchableLicenseMetadata
110+
from datamasque.client.models.rg_config import RGConfig, RGConfigId
108111
from datamasque.client.models.ruleset import Ruleset, RulesetId, RulesetType
109112
from datamasque.client.models.ruleset_library import RulesetLibrary, RulesetLibraryId
110113
from datamasque.client.models.runs import (
@@ -214,6 +217,7 @@
214217
"FileId",
215218
"FileOrContent",
216219
"FileRulesetGenerationRequest",
220+
"FileRulesetGenerationWithRGConfigRequest",
217221
"FirstCharsStatistics",
218222
"ForeignKeyRef",
219223
"GitSnapshot",
@@ -252,9 +256,12 @@
252256
"PatternComposition",
253257
"PatternEntry",
254258
"PatternsStatistics",
259+
"RGConfig",
260+
"RGConfigId",
255261
"ReferencingForeignKey",
256262
"Ruleset",
257263
"RulesetGenerationRequest",
264+
"RulesetGenerationWithRGConfigRequest",
258265
"RulesetId",
259266
"RulesetLibrary",
260267
"RulesetLibraryId",

datamasque/client/discovery.py

Lines changed: 135 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -25,12 +25,15 @@
2525
FileDataDiscoveryRequest,
2626
FileDiscoveryResult,
2727
FileRulesetGenerationRequest,
28+
FileRulesetGenerationWithRGConfigRequest,
2829
RulesetGenerationRequest,
30+
RulesetGenerationWithRGConfigRequest,
2931
SchemaDiscoveryFromConfigRequest,
3032
SchemaDiscoveryPage,
3133
SchemaDiscoveryRequest,
3234
SchemaDiscoveryResult,
3335
)
36+
from datamasque.client.models.rg_config import RGConfig, RGConfigId, unwrap_rg_config_id
3437
from datamasque.client.models.ruleset import Ruleset
3538
from datamasque.client.models.runs import RunId
3639
from datamasque.client.models.status import AsyncRulesetGenerationTaskStatus
@@ -41,21 +44,12 @@
4144
class DiscoveryClient(BaseClient):
4245
"""Schema-discovery and ruleset-generation API methods. Mixed into `DataMasqueClient`."""
4346

44-
def start_async_ruleset_generation(self, connection_id: ConnectionId, selected_data: SelectedData) -> None:
45-
"""
46-
Starts async ruleset generation using the most recent discovery results on the given connection.
47-
48-
If the connection is a database connection, `selected_data` should be of type `SelectedColumns`.
49-
If the connection is a file connection, `selected_data` should be of type `SelectedFileData`.
50-
51-
Generation runs asynchronously on the server.
52-
Poll `get_async_ruleset_generation_task_status` until it returns
53-
`AsyncRulesetGenerationTaskStatus.finished`,
54-
then call `get_generated_rulesets` to retrieve the resulting `Ruleset`.
55-
"""
47+
@staticmethod
48+
def _selected_data_payload(selected_data: SelectedData) -> dict:
49+
"""Build the async-generation request-body fields for a column or file selection, validating its shape."""
5650

5751
if not selected_data:
58-
raise ValueError("`selected_data` is a required argument to `start_async_ruleset_generation`.")
52+
raise ValueError("`selected_data` is a required argument to async ruleset generation.")
5953

6054
data: dict = {}
6155
if isinstance(selected_data, SelectedColumns):
@@ -75,12 +69,54 @@ def start_async_ruleset_generation(self, connection_id: ConnectionId, selected_d
7569
data["selected_data"] = [s.model_dump() for s in selected_data.user_selections]
7670
else:
7771
raise TypeError(
78-
f"The argument `selected_data` to `start_async_ruleset_generation` was of an invalid type, "
72+
f"The `selected_data` argument to async ruleset generation was of an invalid type, "
7973
f"expected `SelectedColumns` or `SelectedFileData`, got {type(selected_data)}."
8074
)
8175

76+
return data
77+
78+
def start_async_ruleset_generation(self, connection_id: ConnectionId, selected_data: SelectedData) -> None:
79+
"""
80+
Starts async ruleset generation using the most recent discovery results on the given connection.
81+
82+
Masks are assigned from the server's default RG config;
83+
use `start_async_ruleset_generation_with_rg_config` to generate with a saved RG config.
84+
85+
If the connection is a database connection, `selected_data` should be of type `SelectedColumns`.
86+
If the connection is a file connection, `selected_data` should be of type `SelectedFileData`.
87+
88+
Generation runs asynchronously on the server.
89+
Poll `get_async_ruleset_generation_task_status` until it returns
90+
`AsyncRulesetGenerationTaskStatus.finished`,
91+
then call `get_generated_rulesets` to retrieve the resulting `Ruleset`.
92+
"""
93+
94+
data = self._selected_data_payload(selected_data)
8295
self.make_request(method="POST", path=f"/api/async-generate-ruleset/{connection_id}/", data=data)
8396

97+
def start_async_ruleset_generation_with_rg_config(
98+
self,
99+
connection_id: ConnectionId,
100+
selected_data: SelectedData,
101+
rg_config: Optional[Union[RGConfigId, RGConfig]],
102+
) -> None:
103+
"""
104+
Starts async ruleset generation with a selected RG config mapping discovered labels to masks.
105+
106+
Like `start_async_ruleset_generation`, but posts to the v2 endpoint with a required `rg_config`:
107+
a saved RG config (`RGConfigId` or `RGConfig`), or `None` for the server's default RG config.
108+
109+
Generation runs asynchronously on the server.
110+
Poll `get_async_ruleset_generation_task_status` until it returns
111+
`AsyncRulesetGenerationTaskStatus.finished`,
112+
then call `get_generated_rulesets` to retrieve the resulting `Ruleset`.
113+
"""
114+
115+
data = self._selected_data_payload(selected_data)
116+
# The server requires `rg_config` to be present; an explicit null selects the default RG config.
117+
data["rg_config"] = unwrap_rg_config_id(rg_config)
118+
self.make_request(method="POST", path=f"/api/async-generate-ruleset/v2/{connection_id}/", data=data)
119+
84120
def start_async_ruleset_generation_from_csv(
85121
self,
86122
connection_id: ConnectionId,
@@ -108,6 +144,52 @@ def start_async_ruleset_generation_from_csv(
108144
then call `get_generated_rulesets` to retrieve the resulting `Ruleset` objects.
109145
"""
110146

147+
self.make_request(
148+
method="POST",
149+
path=f"/api/async-generate-ruleset/{connection_id}/from-csv/",
150+
data={"target_size_bytes": target_size_bytes} if target_size_bytes is not None else None,
151+
files=self._csv_upload_files(csv_content),
152+
)
153+
154+
def start_async_ruleset_generation_from_csv_with_rg_config(
155+
self,
156+
connection_id: ConnectionId,
157+
csv_content: Union[str, bytes, TextIOBase, BufferedIOBase],
158+
rg_config: Optional[Union[RGConfigId, RGConfig]],
159+
target_size_bytes: Optional[int] = None,
160+
) -> None:
161+
"""
162+
Generate ruleset(s) from a schema discovery CSV report with a selected RG config.
163+
164+
Like `start_async_ruleset_generation_from_csv`, but posts to the v2 endpoint with a required
165+
`rg_config`: a saved RG config (`RGConfigId` or `RGConfig`), or `None` for the server's
166+
default RG config.
167+
168+
Generation runs asynchronously on the server.
169+
Poll `get_async_ruleset_generation_task_status` until it returns
170+
`AsyncRulesetGenerationTaskStatus.finished`,
171+
then call `get_generated_rulesets` to retrieve the resulting `Ruleset` objects.
172+
"""
173+
174+
rg_config_id = unwrap_rg_config_id(rg_config)
175+
# The upload is a multipart form, whose fields cannot carry a JSON null;
176+
# the server reads an empty string as null for this nullable field,
177+
# selecting the default RG config.
178+
data: dict = {"rg_config": rg_config_id if rg_config_id is not None else ""}
179+
if target_size_bytes is not None:
180+
data["target_size_bytes"] = target_size_bytes
181+
182+
self.make_request(
183+
method="POST",
184+
path=f"/api/async-generate-ruleset/v2/{connection_id}/from-csv/",
185+
data=data,
186+
files=self._csv_upload_files(csv_content),
187+
)
188+
189+
@staticmethod
190+
def _csv_upload_files(csv_content: Union[str, bytes, TextIOBase, BufferedIOBase]) -> list[UploadFile]:
191+
"""Normalise CSV-or-zip report content into the `csv_or_zip_file` multipart upload."""
192+
111193
content: BufferedIOBase
112194
if isinstance(csv_content, str):
113195
content = BytesIO(csv_content.encode())
@@ -125,7 +207,7 @@ def start_async_ruleset_generation_from_csv(
125207
filename = "ruleset.zip" if is_zip else "ruleset.csv"
126208
content_type = "application/zip" if is_zip else "text/csv"
127209

128-
files = [
210+
return [
129211
UploadFile(
130212
field_name="csv_or_zip_file",
131213
filename=filename,
@@ -134,13 +216,6 @@ def start_async_ruleset_generation_from_csv(
134216
),
135217
]
136218

137-
self.make_request(
138-
method="POST",
139-
path=f"/api/async-generate-ruleset/{connection_id}/from-csv/",
140-
data={"target_size_bytes": target_size_bytes} if target_size_bytes is not None else None,
141-
files=files,
142-
)
143-
144219
def get_async_ruleset_generation_task_status(self, connection_id: ConnectionId) -> AsyncRulesetGenerationTaskStatus:
145220
"""Queries the status of an async ruleset generation task."""
146221

@@ -433,24 +508,62 @@ def generate_ruleset(self, generation_request: RulesetGenerationRequest) -> str:
433508
"""
434509
Generates database-masking ruleset YAML from the most recent discovery run on the given connection.
435510
511+
Masks are assigned from the server's default RG config;
512+
use `generate_ruleset_with_rg_config` to generate with a saved RG config.
513+
436514
`generation_request` is a `RulesetGenerationRequest`.
437515
"""
438516

439517
data = generation_request.model_dump(exclude_none=True, mode="json")
440518
response = self.make_request("POST", "/api/generate-ruleset/v2/", data=data)
441519
return response.content.decode("utf-8")
442520

521+
def generate_ruleset_with_rg_config(self, generation_request: RulesetGenerationWithRGConfigRequest) -> str:
522+
"""
523+
Generates database-masking ruleset YAML with a selected RG config mapping discovered labels to masks.
524+
525+
`generation_request` is a `RulesetGenerationWithRGConfigRequest`;
526+
its required `rg_config` selects the saved RG config to use
527+
(`None` for the server's default RG config).
528+
"""
529+
530+
data = generation_request.model_dump(exclude_none=True, mode="json")
531+
# The server requires `rg_config` to be present; a null selects its default RG config,
532+
# so send it explicitly rather than letting `exclude_none` drop a None.
533+
data.setdefault("rg_config", None)
534+
response = self.make_request("POST", "/api/generate-ruleset/v3/", data=data)
535+
return response.content.decode("utf-8")
536+
443537
def generate_file_ruleset(self, generation_request: FileRulesetGenerationRequest) -> str:
444538
"""
445539
Generates file-masking ruleset YAML from the most recent file-data-discovery run on the given connection.
446540
541+
Masks are assigned from the server's default RG config;
542+
use `generate_file_ruleset_with_rg_config` to generate with a saved RG config.
543+
447544
`generation_request` is a `FileRulesetGenerationRequest`.
448545
"""
449546

450547
data = generation_request.model_dump(exclude_none=True, mode="json")
451548
response = self.make_request("POST", "/api/generate-file-ruleset/", data=data)
452549
return response.content.decode("utf-8")
453550

551+
def generate_file_ruleset_with_rg_config(self, generation_request: FileRulesetGenerationWithRGConfigRequest) -> str:
552+
"""
553+
Generates file-masking ruleset YAML with a selected RG config mapping discovered labels to masks.
554+
555+
`generation_request` is a `FileRulesetGenerationWithRGConfigRequest`;
556+
its required `rg_config` selects the saved RG config to use
557+
(`None` for the server's default RG config).
558+
"""
559+
560+
data = generation_request.model_dump(exclude_none=True, mode="json")
561+
# The server requires `rg_config` to be present; a null selects its default RG config,
562+
# so send it explicitly rather than letting `exclude_none` drop a None.
563+
data.setdefault("rg_config", None)
564+
response = self.make_request("POST", "/api/generate-file-ruleset/v2/", data=data)
565+
return response.content.decode("utf-8")
566+
454567
def get_file_data_discovery_report(self, run_id: RunId) -> list[FileDiscoveryResult]:
455568
"""Returns the file-data-discovery results for the specified run."""
456569

datamasque/client/dmclient.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from datamasque.client.discovery_configs import DiscoveryConfigClient
66
from datamasque.client.files import FileClient
77
from datamasque.client.license import LicenseClient
8+
from datamasque.client.rg_configs import RGConfigClient
89
from datamasque.client.ruleset_libraries import RulesetLibraryClient
910
from datamasque.client.rulesets import RulesetClient
1011
from datamasque.client.runs import RunClient
@@ -26,6 +27,7 @@ class DataMasqueClient(
2627
DiscoveryConfigClient,
2728
DiscoveryConfigLibraryClient,
2829
TableReferenceClient,
30+
RGConfigClient,
2931
UserClient,
3032
SettingsClient,
3133
):

0 commit comments

Comments
 (0)