20
20
21
21
import collections .abc
22
22
import contextlib
23
+ import functools
24
+ import inspect
23
25
import logging
24
26
import re
25
27
import sys
@@ -118,9 +120,11 @@ def mask_secret(secret: str | dict | Iterable, name: str | None = None) -> None:
118
120
_secrets_masker ().add_mask (secret , name )
119
121
120
122
121
- def redact (value : Redactable , name : str | None = None , max_depth : int | None = None ) -> Redacted :
122
- """Redact any secrets found in ``value``."""
123
- return _secrets_masker ().redact (value , name , max_depth )
123
+ def redact (
124
+ value : Redactable , name : str | None = None , max_depth : int | None = None , replacement : str = "***"
125
+ ) -> Redacted :
126
+ """Redact any secrets found in ``value`` with the given replacement."""
127
+ return _secrets_masker ().redact (value , name , max_depth , replacement = replacement )
124
128
125
129
126
130
@overload
@@ -198,6 +202,29 @@ def __init__(self):
198
202
super ().__init__ ()
199
203
self .patterns = set ()
200
204
205
+ @classmethod
206
+ def __init_subclass__ (cls , ** kwargs ):
207
+ super ().__init_subclass__ (** kwargs )
208
+
209
+ if cls ._redact is not SecretsMasker ._redact :
210
+ sig = inspect .signature (cls ._redact )
211
+ # Compat for older versions of the OpenLineage plugin which subclasses this -- call the method
212
+ # without the replacement character
213
+ for param in sig .parameters .values ():
214
+ if param .name == "replacement" or param .kind == param .VAR_KEYWORD :
215
+ break
216
+ else :
217
+ # Block only runs if no break above.
218
+
219
+ f = cls ._redact
220
+
221
+ @functools .wraps (f )
222
+ def _redact (* args , replacement : str = "***" , ** kwargs ):
223
+ return f (* args , ** kwargs )
224
+
225
+ cls ._redact = _redact
226
+ ...
227
+
201
228
@cached_property
202
229
def _record_attrs_to_ignore (self ) -> Iterable [str ]:
203
230
# Doing log.info(..., extra={'foo': 2}) sets extra properties on
@@ -251,59 +278,85 @@ def filter(self, record) -> bool:
251
278
252
279
# Default on `max_depth` is to support versions of the OpenLineage plugin (not the provider) which called
253
280
# this function directly. New versions of that provider, and this class itself call it with a value
254
- def _redact_all (self , item : Redactable , depth : int , max_depth : int = MAX_RECURSION_DEPTH ) -> Redacted :
281
+ def _redact_all (
282
+ self ,
283
+ item : Redactable ,
284
+ depth : int ,
285
+ max_depth : int = MAX_RECURSION_DEPTH ,
286
+ * ,
287
+ replacement : str = "***" ,
288
+ ) -> Redacted :
255
289
if depth > max_depth or isinstance (item , str ):
256
- return "***"
290
+ return replacement
257
291
if isinstance (item , dict ):
258
292
return {
259
- dict_key : self ._redact_all (subval , depth + 1 , max_depth ) for dict_key , subval in item .items ()
293
+ dict_key : self ._redact_all (subval , depth + 1 , max_depth , replacement = replacement )
294
+ for dict_key , subval in item .items ()
260
295
}
261
296
if isinstance (item , (tuple , set )):
262
297
# Turn set in to tuple!
263
- return tuple (self ._redact_all (subval , depth + 1 , max_depth ) for subval in item )
298
+ return tuple (
299
+ self ._redact_all (subval , depth + 1 , max_depth , replacement = replacement ) for subval in item
300
+ )
264
301
if isinstance (item , list ):
265
- return list (self ._redact_all (subval , depth + 1 , max_depth ) for subval in item )
302
+ return list (
303
+ self ._redact_all (subval , depth + 1 , max_depth , replacement = replacement ) for subval in item
304
+ )
266
305
return item
267
306
268
- def _redact (self , item : Redactable , name : str | None , depth : int , max_depth : int ) -> Redacted :
307
+ def _redact (
308
+ self , item : Redactable , name : str | None , depth : int , max_depth : int , replacement : str = "***"
309
+ ) -> Redacted :
269
310
# Avoid spending too much effort on redacting on deeply nested
270
311
# structures. This also avoid infinite recursion if a structure has
271
312
# reference to self.
272
313
if depth > max_depth :
273
314
return item
274
315
try :
275
316
if name and should_hide_value_for_key (name ):
276
- return self ._redact_all (item , depth , max_depth )
317
+ return self ._redact_all (item , depth , max_depth , replacement = replacement )
277
318
if isinstance (item , dict ):
278
319
to_return = {
279
- dict_key : self ._redact (subval , name = dict_key , depth = (depth + 1 ), max_depth = max_depth )
320
+ dict_key : self ._redact (
321
+ subval , name = dict_key , depth = (depth + 1 ), max_depth = max_depth , replacement = replacement
322
+ )
280
323
for dict_key , subval in item .items ()
281
324
}
282
325
return to_return
283
326
if isinstance (item , Enum ):
284
- return self ._redact (item = item .value , name = name , depth = depth , max_depth = max_depth )
327
+ return self ._redact (
328
+ item = item .value , name = name , depth = depth , max_depth = max_depth , replacement = replacement
329
+ )
285
330
if _is_v1_env_var (item ):
286
331
tmp = item .to_dict ()
287
332
if should_hide_value_for_key (tmp .get ("name" , "" )) and "value" in tmp :
288
- tmp ["value" ] = "***"
333
+ tmp ["value" ] = replacement
289
334
else :
290
- return self ._redact (item = tmp , name = name , depth = depth , max_depth = max_depth )
335
+ return self ._redact (
336
+ item = tmp , name = name , depth = depth , max_depth = max_depth , replacement = replacement
337
+ )
291
338
return tmp
292
339
if isinstance (item , str ):
293
340
if self .replacer :
294
341
# We can't replace specific values, but the key-based redacting
295
342
# can still happen, so we can't short-circuit, we need to walk
296
343
# the structure.
297
- return self .replacer .sub ("***" , str (item ))
344
+ return self .replacer .sub (replacement , str (item ))
298
345
return item
299
346
if isinstance (item , (tuple , set )):
300
347
# Turn set in to tuple!
301
348
return tuple (
302
- self ._redact (subval , name = None , depth = (depth + 1 ), max_depth = max_depth ) for subval in item
349
+ self ._redact (
350
+ subval , name = None , depth = (depth + 1 ), max_depth = max_depth , replacement = replacement
351
+ )
352
+ for subval in item
303
353
)
304
354
if isinstance (item , list ):
305
355
return [
306
- self ._redact (subval , name = None , depth = (depth + 1 ), max_depth = max_depth ) for subval in item
356
+ self ._redact (
357
+ subval , name = None , depth = (depth + 1 ), max_depth = max_depth , replacement = replacement
358
+ )
359
+ for subval in item
307
360
]
308
361
return item
309
362
# I think this should never happen, but it does not hurt to leave it just in case
@@ -325,10 +378,12 @@ def _merge(
325
378
self ,
326
379
new_item : Redacted ,
327
380
old_item : Redactable ,
381
+ * ,
328
382
name : str | None ,
329
383
depth : int ,
330
384
max_depth : int ,
331
385
force_sensitive : bool = False ,
386
+ replacement : str ,
332
387
) -> Redacted :
333
388
"""Merge a redacted item with its original unredacted counterpart."""
334
389
if depth > max_depth :
@@ -353,6 +408,7 @@ def _merge(
353
408
depth = depth + 1 ,
354
409
max_depth = max_depth ,
355
410
force_sensitive = is_sensitive ,
411
+ replacement = replacement ,
356
412
)
357
413
else :
358
414
merged [key ] = new_item [key ]
@@ -374,6 +430,7 @@ def _merge(
374
430
depth = depth + 1 ,
375
431
max_depth = max_depth ,
376
432
force_sensitive = is_sensitive ,
433
+ replacement = replacement ,
377
434
)
378
435
)
379
436
else :
@@ -398,25 +455,38 @@ def _merge(
398
455
except (TypeError , AttributeError , ValueError ):
399
456
return new_item
400
457
401
- def redact (self , item : Redactable , name : str | None = None , max_depth : int | None = None ) -> Redacted :
458
+ def redact (
459
+ self ,
460
+ item : Redactable ,
461
+ name : str | None = None ,
462
+ max_depth : int | None = None ,
463
+ replacement : str = "***" ,
464
+ ) -> Redacted :
402
465
"""
403
466
Redact an any secrets found in ``item``, if it is a string.
404
467
405
468
If ``name`` is given, and it's a "sensitive" name (see
406
469
:func:`should_hide_value_for_key`) then all string values in the item
407
470
is redacted.
408
471
"""
409
- return self ._redact (item , name , depth = 0 , max_depth = max_depth or self .MAX_RECURSION_DEPTH )
472
+ return self ._redact (
473
+ item , name , depth = 0 , max_depth = max_depth or self .MAX_RECURSION_DEPTH , replacement = replacement
474
+ )
410
475
411
476
def merge (
412
- self , new_item : Redacted , old_item : Redactable , name : str | None = None , max_depth : int | None = None
477
+ self ,
478
+ new_item : Redacted ,
479
+ old_item : Redactable ,
480
+ name : str | None = None ,
481
+ max_depth : int | None = None ,
482
+ replacement : str = "***" ,
413
483
) -> Redacted :
414
484
"""
415
485
Merge a redacted item with its original unredacted counterpart.
416
486
417
487
Takes a user-modified redacted item and merges it with the original unredacted item.
418
- For sensitive fields that still contain "***" (unchanged), the original value is restored.
419
- For fields that have been updated, the new value is preserved.
488
+ For sensitive fields that still contain "***" (or whatever the ``replacement`` is specified as), the
489
+ original value is restored. For fields that have been updated, the new value is preserved.
420
490
"""
421
491
return self ._merge (
422
492
new_item ,
@@ -425,6 +495,7 @@ def merge(
425
495
depth = 0 ,
426
496
max_depth = max_depth or self .MAX_RECURSION_DEPTH ,
427
497
force_sensitive = False ,
498
+ replacement = replacement ,
428
499
)
429
500
430
501
@cached_property
0 commit comments