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..88187a3db 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( + DataTableWarning.created_by is None + ).delete() + + # add warnings from metastore to the table + for severiy, message in warnings: + DataTableWarning.create( + { + "message": message, + "severity": severiy, + "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