fix: Force flattened record according to provided flattened schema

singer_sdk/helpers/_flattening.py

@@ -361,8 +361,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 +376,9 @@ def flatten_record(
     Returns:
         A flattened version of the record.
     """
+    assert (flattened_schema is not None) or (max_level is not None), "flattened_schema or max_level must be provided"
+    max_level = max_level or 0
+
     return _flatten_record(
         record_node=record,
         flattened_schema=flattened_schema,
@@ -415,7 +418,9 @@ 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, 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,40 @@
+from singer_sdk._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, None, None, AssertionError,
+                 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)
+        print(result)
+        assert expected == result