From 3bbc4669a53526e273b5b717d2ecf028fc5233a3 Mon Sep 17 00:00:00 2001
From: Allison Suarez Miranda <22477579+allisonsuarez@users.noreply.github.com>
Date: Tue, 15 Jun 2021 11:45:52 -0700
Subject: [PATCH] fix: Update common model to align with databuilder model
 (#1232)

* Bump common

Signed-off-by: Allison Suarez Miranda <asuarezmiranda@lyft.com>

* put back chan ges from main

Signed-off-by: Allison Suarez Miranda <asuarezmiranda@lyft.com>

* more changes back

Signed-off-by: Allison Suarez Miranda <asuarezmiranda@lyft.com>
---
 common/amundsen_common/models/feature.py      | 21 ++++++++++++++-----
 .../amundsen_common/models/generation_code.py | 18 ++++++++++++++++
 common/setup.py                               |  2 +-
 3 files changed, 35 insertions(+), 6 deletions(-)
 create mode 100644 common/amundsen_common/models/generation_code.py

diff --git a/common/amundsen_common/models/feature.py b/common/amundsen_common/models/feature.py
index 0a175b56e2..722729aabe 100644
--- a/common/amundsen_common/models/feature.py
+++ b/common/amundsen_common/models/feature.py
@@ -8,7 +8,7 @@
 from amundsen_common.models.user import User
 from amundsen_common.models.badge import Badge
 from amundsen_common.models.tag import Tag
-from amundsen_common.models.table import Column, ProgrammaticDescription, Watermark
+from amundsen_common.models.table import ProgrammaticDescription
 from marshmallow3_annotations.ext.attrs import AttrsSchema
 
 
@@ -38,6 +38,19 @@ class Meta:
         register_as_scheme = True
 
 
+@attr.s(auto_attribs=True, kw_only=True)
+class FeatureWatermark:
+    key: Optional[str]
+    watermark_type: Optional[str]
+    time: str
+
+
+class FeatureWatermarkSchema(AttrsSchema):
+    class Meta:
+        target = FeatureWatermark
+        register_as_scheme = True
+
+
 @attr.s(auto_attribs=True, kw_only=True)
 class Feature:
     key: Optional[str] = attr.ib(default=None)
@@ -51,13 +64,11 @@ class Feature:
     description: Optional[str] = attr.ib(default=None)
     owners: List[User]
     badges: List[Badge]
-    owner_tags: Optional[List[Tag]]  # non editable
-    tags: List[Tag]  # editable
+    tags: List[Tag]
     programmatic_descriptions: List[ProgrammaticDescription]
-    watermarks: List[Watermark]
+    watermarks: List[FeatureWatermark]
     last_updated_timestamp: Optional[int]
     created_timestamp: Optional[int]
-    partition_column: Optional[Column]
 
 
 class FeatureSchema(AttrsSchema):
diff --git a/common/amundsen_common/models/generation_code.py b/common/amundsen_common/models/generation_code.py
new file mode 100644
index 0000000000..ccfdd32ab1
--- /dev/null
+++ b/common/amundsen_common/models/generation_code.py
@@ -0,0 +1,18 @@
+from typing import Optional
+
+import attr
+
+from marshmallow3_annotations.ext.attrs import AttrsSchema
+
+
+@attr.s(auto_attribs=True, kw_only=True)
+class GenerationCode:
+    key: Optional[str]
+    text: str
+    source: Optional[str]
+
+
+class GenerationCodeSchema(AttrsSchema):
+    class Meta:
+        target = GenerationCode
+        register_as_scheme = True
diff --git a/common/setup.py b/common/setup.py
index fc3ccd946e..6630480242 100644
--- a/common/setup.py
+++ b/common/setup.py
@@ -4,7 +4,7 @@
 
 from setuptools import find_packages, setup
 
-__version__ = '0.13.0'
+__version__ = '0.14.0'
 
 requirements_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'requirements-common.txt')
 with open(requirements_path) as requirements_file: