From b39d35811a96c961091d5a6836514355bb3e3601 Mon Sep 17 00:00:00 2001
From: "J.C. Zhong" <jianchengzhong@pinterest.com>
Date: Thu, 24 Aug 2023 17:25:10 -0700
Subject: [PATCH] feat: load table warnings from metastore (#1317)

* feat: load table warnings from metastore

* comments
---
 querybook/server/const/metastore.py           | 13 ++++---
 .../lib/metastore/base_metastore_loader.py    | 11 +++++-
 querybook/server/logic/metastore.py           | 36 +++++++++++++++++--
 3 files changed, 51 insertions(+), 9 deletions(-)

diff --git a/querybook/server/const/metastore.py b/querybook/server/const/metastore.py
index 28cf7ca48..065383436 100644
--- a/querybook/server/const/metastore.py
+++ b/querybook/server/const/metastore.py
@@ -29,6 +29,11 @@ class DataOwner(NamedTuple):
     type: str = None
 
 
+class DataTableWarningSeverity(Enum):
+    WARNING = 0
+    ERROR = 1
+
+
 class DataTable(NamedTuple):
     name: str
 
@@ -73,6 +78,9 @@ class DataTable(NamedTuple):
     golden: bool = False
     boost_score: float = 1
 
+    # table warnings
+    warnings: list[tuple[DataTableWarningSeverity, str]] = None
+
 
 class DataColumn(NamedTuple):
     name: str
@@ -91,11 +99,6 @@ class DataColumn(NamedTuple):
     data_element: DataElementAssociationTuple = None
 
 
-class DataTableWarningSeverity(Enum):
-    WARNING = 0
-    ERROR = 1
-
-
 class MetadataType(Enum):
     TABLE_DESCRIPTION = "table_description"
     COLUMN_DESCRIPTION = "column_description"
diff --git a/querybook/server/lib/metastore/base_metastore_loader.py b/querybook/server/lib/metastore/base_metastore_loader.py
index 3c04e55f0..18d9a89f9 100644
--- a/querybook/server/lib/metastore/base_metastore_loader.py
+++ b/querybook/server/lib/metastore/base_metastore_loader.py
@@ -25,6 +25,7 @@
     create_table,
     create_table_information,
     create_table_ownerships,
+    create_table_warnings,
     delete_column,
     delete_schema,
     delete_table,
@@ -328,7 +329,7 @@ def _create_tables_batched(self, schema_tables):
 
     def _create_tables(self, schema_tables):
         with DBSession() as session:
-            for (schema_id, schema_name, table) in schema_tables:
+            for schema_id, schema_name, table in schema_tables:
                 self._create_table_table(schema_id, schema_name, table, session=session)
 
     @with_session
@@ -380,6 +381,14 @@ def _create_table_table(
                 custom_properties=table.custom_properties,
                 session=session,
             )
+            if table.warnings is not None:
+                create_table_warnings(
+                    table_id=table_id,
+                    warnings=table.warnings,
+                    commit=False,
+                    session=session,
+                )
+
             delete_column_not_in_metastore(
                 table_id, set(map(lambda c: c.name, columns)), session=session
             )
diff --git a/querybook/server/logic/metastore.py b/querybook/server/logic/metastore.py
index f9e484d6d..1f9d92131 100644
--- a/querybook/server/logic/metastore.py
+++ b/querybook/server/logic/metastore.py
@@ -2,7 +2,7 @@
 
 from app.db import with_session
 from const.elasticsearch import ElasticsearchItem
-from const.metastore import DataOwner
+from const.metastore import DataOwner, DataTableWarningSeverity
 from lib.logger import get_logger
 from lib.sqlalchemy import update_model_fields
 from logic.user import get_user_by_name
@@ -17,6 +17,7 @@
     DataTableOwnership,
     DataTableQueryExecution,
     DataTableStatistics,
+    DataTableWarning,
 )
 from models.query_execution import QueryExecution
 from sqlalchemy import and_, func
@@ -75,7 +76,6 @@ def get_schema_by_id(schema_id, session=None):
 
 @with_session
 def update_schema(schema_id, description, session=None):
-
     schema = get_schema_by_id(schema_id, session=session)
 
     schema.description = description
@@ -320,6 +320,37 @@ def create_table_information(
     return table_information
 
 
+@with_session
+def create_table_warnings(
+    table_id,
+    warnings: tuple[DataTableWarningSeverity, str] = [],
+    commit=False,
+    session=None,
+):
+    """This function is used for loading table warnings from metastore.
+
+    For warnings from metastore, created_by will be None.
+    """
+    # delete all warnings without created_by from the table
+    session.query(DataTableWarning).filter_by(
+        table_id=table_id, created_by=None
+    ).delete()
+
+    # add warnings from metastore to the table
+    for severity, message in warnings:
+        DataTableWarning.create(
+            {
+                "message": message,
+                "severity": severity,
+                "table_id": table_id,
+            }
+        )
+    if commit:
+        session.commit()
+    else:
+        session.flush()
+
+
 @with_session
 def delete_table(table_id=None, commit=True, session=None):
     table = get_table_by_id(table_id=table_id, session=session)
@@ -525,7 +556,6 @@ def update_column_by_id(
     commit=True,
     session=None,
 ):
-
     table_column = get_column_by_id(id, session=session)
     if not table_column:
         return