Enable serialization context (#10094) (#10104)

* Update __post_serialize__ signatures

* Temporarily linke dbt-common and dbt-adapters branches

* Changie

* Move fields not in artifacts to resource __post_serialize__ methods

* remove defer_relation in snapshots

* Remove references to branch changes

Co-authored-by: Gerda Shank <[email protected]>

.changes/unreleased/Under the Hood-20240506-145511.yaml

@@ -0,0 +1,6 @@
+kind: Under the Hood
+body: Enable use of context in serialization
+time: 2024-05-06T14:55:11.1812-04:00
+custom:
+  Author: gshank
+  Issue: "10093"

core/dbt/artifacts/resources/v1/components.py

@@ -196,6 +196,12 @@ class ParsedResource(ParsedResourceMandatory):
     relation_name: Optional[str] = None
     raw_code: str = ""
 
+    def __post_serialize__(self, dct: Dict, context: Optional[Dict] = None):
+        dct = super().__post_serialize__(dct, context)
+        if context and context.get("artifact") and "config_call_dict" in dct:
+            del dct["config_call_dict"]
+        return dct
+
 
 @dataclass
 class CompiledResource(ParsedResource):
@@ -215,8 +221,8 @@ class CompiledResource(ParsedResource):
     _pre_injected_sql: Optional[str] = None
     contract: Contract = field(default_factory=Contract)
 
-    def __post_serialize__(self, dct):
-        dct = super().__post_serialize__(dct)
+    def __post_serialize__(self, dct: Dict, context: Optional[Dict] = None):
+        dct = super().__post_serialize__(dct, context)
         if "_pre_injected_sql" in dct:
             del dct["_pre_injected_sql"]
         # Remove compiled attributes

core/dbt/artifacts/resources/v1/model.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Literal, Optional, List
+from typing import Dict, Literal, Optional, List
 from datetime import datetime
 from dbt_common.contracts.config.base import MergeBehavior
 from dbt_common.contracts.constraints import ModelLevelConstraint
@@ -26,3 +26,9 @@ class Model(CompiledResource):
     latest_version: Optional[NodeVersion] = None
     deprecation_date: Optional[datetime] = None
     defer_relation: Optional[DeferRelation] = None
+
+    def __post_serialize__(self, dct: Dict, context: Optional[Dict] = None):
+        dct = super().__post_serialize__(dct, context)
+        if context and context.get("artifact") and "defer_relation" in dct:
+            del dct["defer_relation"]
+        return dct

core/dbt/artifacts/resources/v1/seed.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Optional, Literal
+from typing import Dict, Optional, Literal
 from dbt_common.dataclass_schema import ValidationError
 from dbt.artifacts.resources.types import NodeType
 from dbt.artifacts.resources.v1.components import MacroDependsOn, DeferRelation, ParsedResource
@@ -28,3 +28,9 @@ class Seed(ParsedResource):  # No SQLDefaults!
     root_path: Optional[str] = None
     depends_on: MacroDependsOn = field(default_factory=MacroDependsOn)
     defer_relation: Optional[DeferRelation] = None
+
+    def __post_serialize__(self, dct: Dict, context: Optional[Dict] = None):
+        dct = super().__post_serialize__(dct, context)
+        if context and context.get("artifact") and "defer_relation" in dct:
+            del dct["defer_relation"]
+        return dct

core/dbt/artifacts/resources/v1/snapshot.py

@@ -1,4 +1,4 @@
-from typing import Union, List, Optional, Literal
+from typing import Dict, Union, List, Optional, Literal
 from dataclasses import dataclass
 from dbt_common.dataclass_schema import ValidationError
 from dbt.artifacts.resources.types import NodeType
@@ -64,3 +64,9 @@ class Snapshot(CompiledResource):
     resource_type: Literal[NodeType.Snapshot]
     config: SnapshotConfig
     defer_relation: Optional[DeferRelation] = None
+
+    def __post_serialize__(self, dct, context: Optional[Dict] = None):
+        dct = super().__post_serialize__(dct, context)
+        if context and context.get("artifact") and "defer_relation" in dct:
+            del dct["defer_relation"]
+        return dct

core/dbt/artifacts/schemas/base.py

@@ -34,7 +34,7 @@ def __str__(self) -> str:
 
 class Writable:
     def write(self, path: str):
-        write_json(path, self.to_dict(omit_none=False))  # type: ignore
+        write_json(path, self.to_dict(omit_none=False, context={"artifact": True}))  # type: ignore
 
 
 class Readable:
@@ -60,8 +60,8 @@ class BaseArtifactMetadata(dbtClassMixin):
     invocation_id: Optional[str] = dataclasses.field(default_factory=get_invocation_id)
     env: Dict[str, str] = dataclasses.field(default_factory=get_metadata_vars)
 
-    def __post_serialize__(self, dct):
-        dct = super().__post_serialize__(dct)
+    def __post_serialize__(self, dct: Dict, context: Optional[Dict] = None):
+        dct = super().__post_serialize__(dct, context)
         if dct["generated_at"] and dct["generated_at"].endswith("+00:00"):
             dct["generated_at"] = dct["generated_at"].replace("+00:00", "") + "Z"
         return dct

core/dbt/artifacts/schemas/catalog/v1/catalog.py

@@ -77,8 +77,8 @@ class CatalogResults(dbtClassMixin):
     errors: Optional[List[str]] = None
     _compile_results: Optional[Any] = None
 
-    def __post_serialize__(self, dct):
-        dct = super().__post_serialize__(dct)
+    def __post_serialize__(self, dct: Dict, context: Optional[Dict] = None):
+        dct = super().__post_serialize__(dct, context)
         if "_compile_results" in dct:
             del dct["_compile_results"]
         return dct

core/dbt/artifacts/schemas/manifest/v12/manifest.py

@@ -181,11 +181,3 @@ def upgrade_schema_version(cls, data):
         if manifest_schema_version < cls.dbt_schema_version.version:
             data = upgrade_manifest_json(data, manifest_schema_version)
         return cls.from_dict(data)
-
-    def __post_serialize__(self, dct):
-        for unique_id, node in dct["nodes"].items():
-            if "config_call_dict" in node:
-                del node["config_call_dict"]
-            if "defer_relation" in node:
-                del node["defer_relation"]
-        return dct

core/dbt/contracts/files.py

@@ -139,8 +139,8 @@ def _deserialize(cls, dct: Dict[str, int]):
             sf = SourceFile.from_dict(dct)
         return sf
 
-    def __post_serialize__(self, dct):
-        dct = super().__post_serialize__(dct)
+    def __post_serialize__(self, dct: Dict, context: Optional[Dict] = None):
+        dct = super().__post_serialize__(dct, context)
         # remove empty lists to save space
         dct_keys = list(dct.keys())
         for key in dct_keys:
@@ -226,8 +226,8 @@ def macro_patches(self):
     def source_patches(self):
         return self.sop
 
-    def __post_serialize__(self, dct):
-        dct = super().__post_serialize__(dct)
+    def __post_serialize__(self, dct: Dict, context: Optional[Dict] = None):
+        dct = super().__post_serialize__(dct, context)
         # Remove partial parsing specific data
         for key in ("pp_test_index", "pp_dict"):
             if key in dct:

core/dbt/contracts/graph/manifest.py

@@ -2,7 +2,6 @@
 from collections import defaultdict
 from dataclasses import dataclass, field, replace
 from itertools import chain
-from mashumaro.mixins.msgpack import DataClassMessagePackMixin
 from multiprocessing.synchronize import Lock
 from typing import (
     DefaultDict,
@@ -803,7 +802,7 @@ class ManifestStateCheck(dbtClassMixin):
 
 
 @dataclass
-class Manifest(MacroMethods, DataClassMessagePackMixin, dbtClassMixin):
+class Manifest(MacroMethods, dbtClassMixin):
     """The manifest for the full graph, after parsing and during compilation."""
 
     # These attributes are both positional and by keyword. If an attribute
@@ -870,7 +869,7 @@ class Manifest(MacroMethods, DataClassMessagePackMixin, dbtClassMixin):
         metadata={"serialize": lambda x: None, "deserialize": lambda x: None},
     )
 
-    def __pre_serialize__(self):
+    def __pre_serialize__(self, context: Optional[Dict] = None):
         # serialization won't work with anything except an empty source_patches because
         # tuple keys are not supported, so ensure it's empty
         self.source_patches = {}

core/dbt/contracts/graph/nodes.py

@@ -262,8 +262,8 @@ def write_node(self, project_root: str, compiled_path, compiled_code: str):
     def _serialize(self):
         return self.to_dict()
 
-    def __post_serialize__(self, dct):
-        dct = super().__post_serialize__(dct)
+    def __post_serialize__(self, dct: Dict, context: Optional[Dict] = None):
+        dct = super().__post_serialize__(dct, context)
         if "_event_status" in dct:
             del dct["_event_status"]
         return dct

core/dbt/contracts/graph/unparsed.py

@@ -276,8 +276,8 @@ class UnparsedSourceTableDefinition(HasColumnTests, HasColumnAndTestProps):
     external: Optional[ExternalTable] = None
     tags: List[str] = field(default_factory=list)
 
-    def __post_serialize__(self, dct):
-        dct = super().__post_serialize__(dct)
+    def __post_serialize__(self, dct: Dict, context: Optional[Dict] = None):
+        dct = super().__post_serialize__(dct, context)
         if "freshness" not in dct and self.freshness is None:
             dct["freshness"] = None
         return dct
@@ -314,8 +314,8 @@ def validate(cls, data):
     def yaml_key(self) -> "str":
         return "sources"
 
-    def __post_serialize__(self, dct):
-        dct = super().__post_serialize__(dct)
+    def __post_serialize__(self, dct: Dict, context: Optional[Dict] = None):
+        dct = super().__post_serialize__(dct, context)
         if "freshness" not in dct and self.freshness is None:
             dct["freshness"] = None
         return dct

core/dbt/parser/manifest.py

@@ -233,7 +233,7 @@ class ManifestLoaderInfo(dbtClassMixin, Writable):
     projects: List[ProjectLoaderInfo] = field(default_factory=list)
     _project_index: Dict[str, ProjectLoaderInfo] = field(default_factory=dict)
 
-    def __post_serialize__(self, dct):
+    def __post_serialize__(self, dct: Dict, context: Optional[Dict] = None):
         del dct["_project_index"]
         return dct
 

third-party-stubs/mashumaro/jsonschema/builder.pyi

@@ -16,7 +16,7 @@ def build_json_schema(
 
 class JSONSchemaDefinitions(DataClassJSONMixin):
     definitions: Dict[str, JSONSchema]
-    def __post_serialize__(self, d: Dict[Any, Any]) -> List[Dict[str, Any]]: ...  # type: ignore
+    def __post_serialize__(self, d: Dict[Any, Any], context: Optional[Dict]) -> List[Dict[str, Any]]: ...  # type: ignore
     def __init__(self, definitions) -> None: ...
 
 class JSONSchemaBuilder:

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,

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]: ...