From 08efc20e36f32f73dcbdceb7f8244ca5f34c42e1 Mon Sep 17 00:00:00 2001
From: matanmichaely <matan.michaely@yotpo.com>
Date: Thu, 19 Sep 2024 11:58:03 +0300
Subject: [PATCH] update models modules

---
 yoda_dbt2looker/core/models.py | 22 ++++++++++++++++++----
 yoda_dbt2looker/models.py      |  2 --
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/yoda_dbt2looker/core/models.py b/yoda_dbt2looker/core/models.py
index 652c18e..6be7918 100644
--- a/yoda_dbt2looker/core/models.py
+++ b/yoda_dbt2looker/core/models.py
@@ -1,10 +1,10 @@
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Union, Any
 
 try:
     from typing import Literal
 except ImportError:
     from typing_extensions import Literal
-from pydantic import Field, validator, BaseModel
+from pydantic import Field, field_validator, BaseModel, ValidationError, model_validator
 
 from yoda_dbt2looker.models import (
     DbtNode,
@@ -34,7 +34,7 @@ class DbtModel(DbtNode):
     meta: DbtModelMeta
     config: DbtModelConfig
 
-    @validator("columns")
+    @field_validator("columns")
     def case_insensitive_column_names(cls, v: Dict[str, DbtModelColumn]):
         return {
             name.lower(): column.copy(update={"name": column.name.lower()})
@@ -43,5 +43,19 @@ def case_insensitive_column_names(cls, v: Dict[str, DbtModelColumn]):
 
 
 class DbtManifest(BaseModel):
-    nodes: Dict[str, Union[DbtModel, DbtNode]]
+    # TODO: check why Union doesn't work
+    nodes: Dict[str, Any]
     metadata: DbtManifestMetadata
+
+    @model_validator(mode="before")
+    def validate_nodes_and_exposures(cls, values):
+        nodes = values.get('nodes', {})
+        # Validate 'nodes' field
+        validated_nodes = {}
+        for key, value in nodes.items():
+            try:
+                validated_nodes[key] = DbtModel(**value)
+            except ValidationError as e:
+                validated_nodes[key] = DbtNode(**value)
+        values['nodes'] = validated_nodes
+        return values
diff --git a/yoda_dbt2looker/models.py b/yoda_dbt2looker/models.py
index 17f891e..61c3fb6 100644
--- a/yoda_dbt2looker/models.py
+++ b/yoda_dbt2looker/models.py
@@ -381,8 +381,6 @@ def validate_nodes_and_exposures(cls, values):
             try:
                 validated_nodes[key] = DbtModel(**value)
             except ValidationError as e:
-                if "model.yoda.yoda_e2e_loyalty__fact_store_profile_daily" == key:
-                    print(f"failed to parse {key} , errors: {e.errors()}")
                 validated_nodes[key] = DbtNode(**value)
 
         # Validate 'exposures' field