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
3437from datamasque .client .models .ruleset import Ruleset
3538from datamasque .client .models .runs import RunId
3639from datamasque .client .models .status import AsyncRulesetGenerationTaskStatus
4144class 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
0 commit comments