fix: Force flattened record according to provided flattened schema

singer_sdk/exceptions.py

@@ -157,3 +157,7 @@ def __init__(self, error_message: str, record: dict) -> None:
         super().__init__(f"Record Message Validation Error: {error_message}")
         self.error_message = error_message
         self.record = record
+
+
+class InvalidFlatteningRecordsParameter(Exception):
+    """Raised when the flattening_records parameter is invalid."""

singer_sdk/helpers/_flattening.py

@@ -11,6 +11,8 @@
 import inflection
 import simplejson as json
 
+from singer_sdk.exceptions import InvalidFlatteningRecordsParameter
+
 DEFAULT_FLATTENING_SEPARATOR = "__"
 
 
@@ -361,8 +363,8 @@ def _key_func(item: tuple[str, dict]) -> str:
 
 def flatten_record(
     record: dict,
-    flattened_schema: dict,
-    max_level: int,
+    flattened_schema: dict | None = None,
+    max_level: int | None = None,
     separator: str = "__",
 ) -> dict:
     """Flatten a record up to max_level.
@@ -376,6 +378,11 @@ def flatten_record(
     Returns:
         A flattened version of the record.
     """
+    if flattened_schema is None and max_level is None:
+        msg = "flattened_schema or max_level must be provided"
+        raise InvalidFlatteningRecordsParameter(msg)
+    max_level = max_level or 0
+
     return _flatten_record(
         record_node=record,
         flattened_schema=flattened_schema,
@@ -415,7 +422,10 @@ def _flatten_record(
     items: list[tuple[str, t.Any]] = []
     for k, v in record_node.items():
         new_key = flatten_key(k, parent_key, separator)
-        if isinstance(v, collections.abc.MutableMapping) and level < max_level:
+        if isinstance(v, collections.abc.MutableMapping) and (
+            (flattened_schema and new_key not in flattened_schema.get("properties", {}))
+            or (not flattened_schema and level < max_level)
+        ):
             items.extend(
                 _flatten_record(
                     v,

tests/core/test_flattening.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import pytest
+
+from singer_sdk.exceptions import InvalidFlatteningRecordsParameter
+from singer_sdk.helpers._flattening import flatten_record
+
+
+@pytest.mark.parametrize(
+    "flattened_schema, max_level, expected, expected_exception",
+    [
+        pytest.param(
+            {
+                "properties": {
+                    "key_1": {"type": ["null", "integer"]},
+                    "key_2__key_3": {"type": ["null", "string"]},
+                    "key_2__key_4": {"type": ["null", "object"]},
+                }
+            },
+            None,
+            {
+                "key_1": 1,
+                "key_2__key_3": "value",
+                "key_2__key_4": '{"key_5": 1, "key_6": ["a", "b"]}',
+            },
+            None,
+            id="flattened schema provided",
+        ),
+        pytest.param(
+            None,
+            99,
+            {
+                "key_1": 1,
+                "key_2__key_3": "value",
+                "key_2__key_4__key_5": 1,
+                "key_2__key_4__key_6": '["a", "b"]',
+            },
+            None,
+            id="flattened schema not provided",
+        ),
+        pytest.param(
+            None,
+            1,
+            {
+                "key_1": 1,
+                "key_2__key_3": "value",
+                "key_2__key_4": '{"key_5": 1, "key_6": ["a", "b"]}',
+            },
+            None,
+            id="limited by max_level 2",
+        ),
+        pytest.param(
+            None,
+            None,
+            None,
+            InvalidFlatteningRecordsParameter,
+            id="no schema or max level provided",
+        ),
+    ],
+)
+def test_flatten_record(flattened_schema, max_level, expected, expected_exception):
+    """Test flatten_record to obey the max_level and flattened_schema parameters."""
+    record = {
+        "key_1": 1,
+        "key_2": {"key_3": "value", "key_4": {"key_5": 1, "key_6": ["a", "b"]}},
+    }
+    if expected_exception:
+        with pytest.raises(expected_exception):
+            flatten_record(
+                record, max_level=max_level, flattened_schema=flattened_schema
+            )
+    else:
+        result = flatten_record(
+            record, max_level=max_level, flattened_schema=flattened_schema
+        )
+        assert expected == result