feat(ingestion): Add secret masking framework (#15188)

Author: kyungsoo-datahub

metadata-ingestion/src/datahub/cli/ingest_cli.py

@@ -21,6 +21,7 @@
 from datahub.ingestion.graph.config import ClientMode
 from datahub.ingestion.run.connection import ConnectionManager
 from datahub.ingestion.run.pipeline import Pipeline
+from datahub.masking.bootstrap import initialize_secret_masking
 from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
 from datahub.utilities.ingest_utils import deploy_source_vars
@@ -128,6 +129,9 @@ def run(
 ) -> None:
     """Ingest metadata into DataHub."""
 
+    # Initialize secret masking (before any logging)
+    initialize_secret_masking()
+
     def run_pipeline_to_completion(pipeline: Pipeline) -> int:
         logger.info("Starting metadata ingestion")
         with click_spinner.spinner(disable=no_spinner or no_progress):

metadata-ingestion/src/datahub/configuration/common.py

@@ -1,3 +1,4 @@
+import contextvars
 import dataclasses
 import re
 import unittest.mock
@@ -21,11 +22,12 @@
 import pydantic
 import pydantic_core
 from cached_property import cached_property
-from pydantic import BaseModel, ConfigDict, ValidationError
+from pydantic import BaseModel, ConfigDict, SecretStr, ValidationError, model_validator
 from pydantic.fields import Field
 from typing_extensions import Protocol, Self
 
 from datahub.configuration._config_enum import ConfigEnum as ConfigEnum
+from datahub.masking.secret_registry import SecretRegistry, is_masking_enabled
 from datahub.utilities.dedup_list import deduplicate_list
 
 REDACT_KEYS = {
@@ -95,6 +97,11 @@ def redact_raw_config(obj: Any) -> Any:
 
 LaxStr = Annotated[str, pydantic.BeforeValidator(lambda v: str(v))]
 
+# Context variable to track if we're inside a nested ConfigModel construction
+_inside_nested_config: contextvars.ContextVar[bool] = contextvars.ContextVar(
+    "_inside_nested_config", default=False
+)
+
 
 @dataclasses.dataclass(frozen=True)
 class SupportedSources:
@@ -129,6 +136,96 @@ class ConfigModel(BaseModel):
         json_schema_extra=_config_model_schema_extra,
     )
 
+    @model_validator(mode="wrap")
+    @classmethod
+    def _track_nesting_context(
+        cls,
+        data: Any,
+        handler: pydantic.ValidatorFunctionWrapHandler,
+        info: pydantic.ValidationInfo,
+    ) -> Self:
+        """
+        Wrap validator that tracks nesting context for nested ConfigModel detection.
+
+        Sets a context variable so nested ConfigModels know they're being constructed as fields.
+        """
+        # Set context for any nested models that will be created during field processing
+        token = _inside_nested_config.set(True)
+        try:
+            # Process the model normally (this calls __init__ and all validators)
+            instance = handler(data)
+        finally:
+            # Reset context after processing
+            _inside_nested_config.reset(token)
+
+        return instance
+
+    @model_validator(mode="after")
+    def _register_secret_fields(self) -> Self:
+        """
+        Register SecretStr fields with the secret masking registry.
+        Recursively traverses nested ConfigModel instances to find all SecretStr fields.
+
+        Only models that are constructed outside of Pydantic field processing will register secrets.
+        This ensures we capture the full qualified paths for nested secrets.
+
+        Performance: Uses batch registration for efficiency - single version
+        increment instead of one per secret.
+        """
+        if not is_masking_enabled():
+            return self
+
+        # Only register if we're NOT inside another ConfigModel's field processing
+        # This means we're a "root" model from the user's perspective
+        if _inside_nested_config.get():
+            return self
+
+        # Collect all secrets recursively (including from nested models)
+        secrets: Dict[str, str] = {}
+        self._collect_secrets(secrets, prefix="")
+
+        # Batch register all secrets in one operation
+        if secrets:
+            SecretRegistry.get_instance().register_secrets_batch(secrets)
+
+        return self
+
+    def _collect_secrets(self, secrets: Dict[str, str], prefix: str) -> None:
+        """
+        Recursively collect SecretStr fields from this model and nested ConfigModel instances.
+
+        Args:
+            secrets: Dictionary to populate with field_name -> secret_value mappings
+            prefix: Prefix for nested field names (e.g., "azure_auth." for nested fields)
+        """
+        for field_name, _field_info in self.__class__.model_fields.items():
+            field_value = getattr(self, field_name, None)
+
+            if field_value is None:
+                continue
+
+            # Build the full field path for better debugging
+            full_name = f"{prefix}{field_name}" if prefix else field_name
+
+            if isinstance(field_value, SecretStr):
+                # Direct SecretStr field
+                secret_value = field_value.get_secret_value()
+                if secret_value:
+                    secrets[full_name] = secret_value
+            elif isinstance(field_value, ConfigModel):
+                # Nested ConfigModel - recurse into it
+                field_value._collect_secrets(secrets, prefix=f"{full_name}.")
+            elif isinstance(field_value, list):
+                # Handle lists of ConfigModels
+                for idx, item in enumerate(field_value):
+                    if isinstance(item, ConfigModel):
+                        item._collect_secrets(secrets, prefix=f"{full_name}[{idx}].")
+            elif isinstance(field_value, dict):
+                # Handle dicts with ConfigModel values
+                for key, item in field_value.items():
+                    if isinstance(item, ConfigModel):
+                        item._collect_secrets(secrets, prefix=f"{full_name}[{key}].")
+
     @classmethod
     def parse_obj_allow_extras(cls, obj: Any) -> Self:
         """Parse an object while allowing extra fields.

metadata-ingestion/src/datahub/configuration/config_loader.py

@@ -15,10 +15,32 @@
 from datahub.configuration.json_loader import JsonConfigurationMechanism
 from datahub.configuration.toml import TomlConfigurationMechanism
 from datahub.configuration.yaml import YamlConfigurationMechanism
+from datahub.masking.secret_registry import SecretRegistry, is_masking_enabled
 
 Environ = Mapping[str, str]
 
 
+def _extract_env_var_names(text: str) -> Set[str]:
+    """Extract environment variable names from ${VAR} or $VAR patterns."""
+    var_names = set()
+
+    # Match ${VAR} and bash parameter expansion patterns
+    # Pattern breakdown:
+    #   ([A-Za-z_][A-Za-z0-9_]*)  - capture variable name
+    #   (?::[+\-=?][^}]*)?        - optional bash operators with content
+    for match in re.finditer(r"\$\{([A-Za-z_][A-Za-z0-9_]*)(?::[+\-=?][^}]*)?\}", text):
+        var_names.add(match.group(1))
+
+    # Match $VAR patterns (without braces)
+    for match in re.finditer(r"\$([A-Za-z_][A-Za-z0-9_]*)", text):
+        var_name = match.group(1)
+        # Only add if not already captured by ${} pattern
+        if var_name not in var_names:
+            var_names.add(var_name)
+
+    return var_names
+
+
 def resolve_env_variables(config: dict, environ: Environ) -> dict:
     # TODO: This is kept around for backwards compatibility.
     return EnvResolver(environ).resolve(config)
@@ -30,9 +52,26 @@ def list_referenced_env_variables(config: dict) -> Set[str]:
 
 
 class EnvResolver:
-    def __init__(self, environ: Environ, strict_env_syntax: bool = False):
+    """Resolves environment variable references in configuration dictionaries."""
+
+    def __init__(
+        self,
+        environ: Environ,
+        strict_env_syntax: bool = False,
+        register_secrets: bool = True,
+    ):
+        """
+        Initialize the environment variable resolver.
+
+        Args:
+            environ: Environment variable mapping (os.environ, custom dict, or external secrets)
+            strict_env_syntax: If True, only match ${VAR} syntax (not $VAR)
+            register_secrets: If True, register resolved values with masking registry
+                            (only if masking is globally enabled)
+        """
         self.environ = environ
         self.strict_env_syntax = strict_env_syntax
+        self.register_secrets = register_secrets
 
     def resolve(self, config: dict) -> dict:
         return self._resolve_dict(config)
@@ -57,16 +96,42 @@ def mock_get_env(key: str, default: Optional[str] = None) -> str:
         mock = unittest.mock.MagicMock()
         mock.get.side_effect = mock_get_env
 
-        resolver = EnvResolver(environ=mock, strict_env_syntax=strict_env_syntax)
+        resolver = EnvResolver(
+            environ=mock, strict_env_syntax=strict_env_syntax, register_secrets=False
+        )
         resolver._resolve_dict(config)
 
         return vars
 
+    def _register_env_vars_from_element(self, element: str) -> None:
+        """Register environment variables found in element for secret masking."""
+        if not self.register_secrets or not is_masking_enabled():
+            return
+
+        # Extract variable names from the pattern
+        var_names = _extract_env_var_names(element)
+
+        # Collect secrets for batch registration
+        # We register all ${VAR} references from recipe as secrets
+        secrets = {}
+        for var_name in var_names:
+            value = self.environ.get(var_name)
+            if value:
+                secrets[var_name] = value
+
+        # Batch register all secrets from this element
+        if secrets:
+            SecretRegistry.get_instance().register_secrets_batch(secrets)
+
     def _resolve_element(self, element: str) -> str:
         if re.search(r"(\$\{).+(\})", element):
+            # Register secrets before expansion
+            self._register_env_vars_from_element(element)
             return expand(element, nounset=True, environ=self.environ)
         elif not self.strict_env_syntax and element.startswith("$"):
             try:
+                # Register secrets before expansion
+                self._register_env_vars_from_element(element)
                 return expand(element, nounset=True, environ=self.environ)
             except UnboundVariable:
                 # TODO: This fallback is kept around for backwards compatibility, but

metadata-ingestion/src/datahub/configuration/env_vars.py

@@ -151,6 +151,16 @@ def get_debug() -> bool:
     return os.getenv("DATAHUB_DEBUG", "").lower() == "true"
 
 
+def get_disable_secret_masking() -> bool:
+    """
+    Disable secret masking for debugging purposes.
+
+    WARNING: Only use this in development/debugging scenarios.
+    Disabling secret masking will expose sensitive information in logs.
+    """
+    return os.getenv("DATAHUB_DISABLE_SECRET_MASKING", "").lower() in ("true", "1")
+
+
 # ============================================================================
 # Data Processing Configuration
 # ============================================================================

metadata-ingestion/src/datahub/masking/__init__.py

@@ -0,0 +1,24 @@
+"""
+Secret masking system for DataHub ingestion.
+"""
+
+__all__ = [
+    "SecretMaskingFilter",
+    "StreamMaskingWrapper",
+    "install_masking_filter",
+    "uninstall_masking_filter",
+    "SecretRegistry",
+    "is_masking_enabled",
+    "initialize_secret_masking",
+    "get_masking_safe_logger",
+]
+
+from datahub.masking.bootstrap import initialize_secret_masking
+from datahub.masking.logging_utils import get_masking_safe_logger
+from datahub.masking.masking_filter import (
+    SecretMaskingFilter,
+    StreamMaskingWrapper,
+    install_masking_filter,
+    uninstall_masking_filter,
+)
+from datahub.masking.secret_registry import SecretRegistry, is_masking_enabled