From f89a5168534f912bdb664c1189f1b25e54033f0c Mon Sep 17 00:00:00 2001
From: Gerda Shank <gerda@dbtlabs.com>
Date: Tue, 7 May 2024 16:25:15 -0400
Subject: [PATCH] Add msgpack mixin to dbtClassMixin (#125)

* Add msgpack mixin to dbtClassMixin

* missing Optional

* Changie
---
 .../unreleased/Features-20240506-170058.yaml  |  6 +++
 dbt_common/contracts/config/properties.py     | 10 ++---
 dbt_common/dataclass_schema.py                | 16 +++++--
 tests/unit/test_dataclass_schema.py           | 42 +++++++++++++++++++
 .../mashumaro/jsonschema/models.pyi           |  4 +-
 third-party-stubs/mashumaro/mixins/dict.pyi   |  6 +--
 6 files changed, 70 insertions(+), 14 deletions(-)
 create mode 100644 .changes/unreleased/Features-20240506-170058.yaml
 create mode 100644 tests/unit/test_dataclass_schema.py

diff --git a/.changes/unreleased/Features-20240506-170058.yaml b/.changes/unreleased/Features-20240506-170058.yaml
new file mode 100644
index 00000000..d378e228
--- /dev/null
+++ b/.changes/unreleased/Features-20240506-170058.yaml
@@ -0,0 +1,6 @@
+kind: Features
+body: Enable serialization contexts
+time: 2024-05-06T17:00:58.272104-04:00
+custom:
+  Author: gshank
+  Issue: "126"
diff --git a/dbt_common/contracts/config/properties.py b/dbt_common/contracts/config/properties.py
index ce669623..5277c86e 100644
--- a/dbt_common/contracts/config/properties.py
+++ b/dbt_common/contracts/config/properties.py
@@ -1,10 +1,10 @@
 from dataclasses import dataclass, field
-from typing import Dict, Any
+from typing import Dict, Any, Optional
 
-from dbt_common.dataclass_schema import ExtensibleDbtClassMixin
+from dbt_common.dataclass_schema import ExtensibleDbtClassMixin, dbtClassMixin
 
 
-class AdditionalPropertiesMixin:
+class AdditionalPropertiesMixin(dbtClassMixin):
     """Make this class an extensible property.
 
     The underlying class definition must include a type definition for a field
@@ -41,8 +41,8 @@ def __pre_deserialize__(cls, data):
         data = super().__pre_deserialize__(data)
         return data
 
-    def __post_serialize__(self, dct):
-        data = super().__post_serialize__(dct)
+    def __post_serialize__(self, dct: Dict, context: Optional[Dict] = None):
+        data = super().__post_serialize__(dct, context)
         data.update(self.extra)
         if "_extra" in data:
             del data["_extra"]
diff --git a/dbt_common/dataclass_schema.py b/dbt_common/dataclass_schema.py
index 9c171cbd..0bad081f 100644
--- a/dbt_common/dataclass_schema.py
+++ b/dbt_common/dataclass_schema.py
@@ -7,11 +7,17 @@
 from dateutil.parser import parse
 
 # type: ignore
-from mashumaro import DataClassDictMixin
-from mashumaro.config import TO_DICT_ADD_OMIT_NONE_FLAG, BaseConfig as MashBaseConfig
+from mashumaro.config import (
+    TO_DICT_ADD_OMIT_NONE_FLAG,
+    ADD_SERIALIZATION_CONTEXT,
+    BaseConfig as MashBaseConfig,
+)
 from mashumaro.types import SerializableType, SerializationStrategy
 from mashumaro.jsonschema import build_json_schema
 
+# following includes DataClassDictMixin
+from mashumaro.mixins.msgpack import DataClassMessagePackMixin
+
 import functools
 
 
@@ -34,6 +40,7 @@ def deserialize(self, value) -> datetime:
 class dbtMashConfig(MashBaseConfig):
     code_generation_options = [
         TO_DICT_ADD_OMIT_NONE_FLAG,
+        ADD_SERIALIZATION_CONTEXT,
     ]
     serialization_strategy = {
         datetime: DateTimeSerialization(),
@@ -47,7 +54,8 @@ class dbtMashConfig(MashBaseConfig):
 
 # This class pulls in DataClassDictMixin from Mashumaro. The 'to_dict'
 # and 'from_dict' methods come from Mashumaro.
-class dbtClassMixin(DataClassDictMixin):
+# Note: DataClassMessagePackMixin inherits from DataClassDictMixin
+class dbtClassMixin(DataClassMessagePackMixin):
     """Convert and validate JSON schemas.
 
     The Mixin adds methods to generate a JSON schema and
@@ -73,7 +81,7 @@ def __pre_deserialize__(cls, data):
     # This is called by the mashumaro to_dict in order to handle
     # nested classes. We no longer do any munging here, but leaving here
     # so that subclasses can leave super() in place for possible future needs.
-    def __post_serialize__(self, data):
+    def __post_serialize__(self, data, context: Optional[Dict]):
         return data
 
     @classmethod
diff --git a/tests/unit/test_dataclass_schema.py b/tests/unit/test_dataclass_schema.py
new file mode 100644
index 00000000..16e50c38
--- /dev/null
+++ b/tests/unit/test_dataclass_schema.py
@@ -0,0 +1,42 @@
+from dataclasses import dataclass
+from typing import Dict, Optional
+
+from dbt_common.dataclass_schema import dbtClassMixin
+
+
+@dataclass
+class MySubObject(dbtClassMixin):
+    name: str
+
+    def __post_serialize__(self, dct: Dict, context: Optional[Dict] = None):
+        if context and "artifact" in context:
+            dct["name"] = "xxxx"
+        return dct
+
+
+@dataclass
+class MyObject(dbtClassMixin):
+    sub_object: MySubObject
+    unique_id: str
+
+    def __post_serialize__(self, dct: Dict, context: Optional[Dict] = None):
+        if context and "artifact" in context:
+            dct["unique_id"] = "my.xxxx.object"
+        return dct
+
+
+def test_serialization_context():
+    sub_obj = MySubObject("testing")
+
+    obj = MyObject(sub_object=sub_obj, unique_id="my.test.object")
+
+    assert obj
+
+    dct = obj.to_dict()
+    assert dct == {"unique_id": "my.test.object", "sub_object": {"name": "testing"}}
+
+    mod_dct = obj.to_dict(context={"artifact": True})
+    assert mod_dct == {"unique_id": "my.xxxx.object", "sub_object": {"name": "xxxx"}}
+
+    obj = MyObject.from_dict(dct)
+    assert obj.sub_object.name == "testing"
diff --git a/third-party-stubs/mashumaro/jsonschema/models.pyi b/third-party-stubs/mashumaro/jsonschema/models.pyi
index b67db67b..6022d3d1 100644
--- a/third-party-stubs/mashumaro/jsonschema/models.pyi
+++ b/third-party-stubs/mashumaro/jsonschema/models.pyi
@@ -106,8 +106,8 @@ class JSONSchema(DataClassJSONMixin):
         serialize_by_alias: bool
         aliases: Incomplete
         serialization_strategy: Incomplete
-    def __pre_serialize__(self) -> JSONSchema: ...
-    def __post_serialize__(self, d: Dict[Any, Any]) -> Dict[Any, Any]: ...
+    def __pre_serialize__(self, context: Optional[Dict]) -> JSONSchema: ...
+    def __post_serialize__(self, d: Dict[Any, Any], context: Optional[Dict]) -> Dict[Any, Any]: ...
     def __init__(
         self,
         schema,
diff --git a/third-party-stubs/mashumaro/mixins/dict.pyi b/third-party-stubs/mashumaro/mixins/dict.pyi
index 87728396..c6ec9acc 100644
--- a/third-party-stubs/mashumaro/mixins/dict.pyi
+++ b/third-party-stubs/mashumaro/mixins/dict.pyi
@@ -1,4 +1,4 @@
-from typing import Any, Dict, Mapping, Type, TypeVar
+from typing import Any, Dict, Mapping, Type, TypeVar, Optional
 
 T = TypeVar("T", bound="DataClassDictMixin")
 
@@ -11,5 +11,5 @@ class DataClassDictMixin:
     def __pre_deserialize__(cls: Type[T], d: Dict[Any, Any]) -> Dict[Any, Any]: ...
     @classmethod
     def __post_deserialize__(cls: Type[T], obj: T) -> T: ...
-    def __pre_serialize__(self: T) -> T: ...
-    def __post_serialize__(self, d: Dict[Any, Any]) -> Dict[Any, Any]: ...
+    def __pre_serialize__(self: T, context: Optional[Dict]) -> T: ...
+    def __post_serialize__(self, d: Dict[Any, Any], context: Optional[Dict]) -> Dict[Any, Any]: ...