diff --git a/querybook/server/const/data_element.py b/querybook/server/const/data_element.py index 9de59af79..c277bdcd0 100644 --- a/querybook/server/const/data_element.py +++ b/querybook/server/const/data_element.py @@ -51,11 +51,11 @@ class DataElementTuple(NamedTuple): class DataElementAssociationTuple(NamedTuple): # association type type: DataElementAssociationType - # data element name. required for all association types - value: str + # data element tuple or name. required for all association types + value_data_element: Union[DataElementTuple, str] # for map association, value can either be a data element or a primitive type, e.g. i32 value_primitive_type: str = None - # data element name. required if associtaion type is map - key: str = None + # data element tuple or name. required if associtaion type is map + key_data_element: Union[DataElementTuple, str] = None # for map association, key can either be a data element or a primitive type, e.g. i32 key_primitive_type: str = None diff --git a/querybook/server/lib/metastore/base_metastore_loader.py b/querybook/server/lib/metastore/base_metastore_loader.py index a6c224173..02b88bc90 100644 --- a/querybook/server/lib/metastore/base_metastore_loader.py +++ b/querybook/server/lib/metastore/base_metastore_loader.py @@ -5,6 +5,7 @@ import gevent from app.db import DBSession, with_session +from const.data_element import DataElementTuple, DataElementAssociationTuple from const.metastore import ( DataColumn, DataOwnerType, @@ -16,6 +17,7 @@ from lib.logger import get_logger from lib.utils import json from lib.utils.utils import with_exception +from logic.data_element import create_column_data_element_association from logic.elasticsearch import delete_es_table_by_id, update_table_by_id from logic.metastore import ( create_column, @@ -33,7 +35,6 @@ iterate_data_schema, ) from logic.tag import create_column_tags, create_table_tags -from logic.data_element import create_column_data_element_association from .utils import MetastoreTableACLChecker @@ -405,9 +406,15 @@ def _create_table_table( if self.loader_config.can_load_external_metadata( MetadataType.DATA_ELEMENT ): + data_element_association = ( + self._populate_column_data_element_association( + column.data_element + ) + ) create_column_data_element_association( + metastore_id=self.metastore_id, column_id=column_id, - data_element_association=column.data_element, + data_element_association=data_element_association, commit=False, session=session, ) @@ -436,6 +443,26 @@ def _create_table_table( session.rollback() LOG.error(traceback.format_exc()) + def _populate_column_data_element_association( + self, data_element_association: DataElementAssociationTuple + ) -> DataElementAssociationTuple: + """If the value_data_element or key_data_element is a name instead of DataElementTuple, + this function will help to replace the name with the actual DataElementTuple""" + if data_element_association is None: + return None + + value_data_element = data_element_association.value_data_element + if type(value_data_element) is str: + value_data_element = self.get_data_element(value_data_element) + + key_data_element = data_element_association.key_data_element + if type(key_data_element) is str: + key_data_element = self.get_data_element(key_data_element) + + return data_element_association._replace( + value_data_element=value_data_element, key_data_element=key_data_element + ) + @with_exception def _get_all_filtered_schema_names(self) -> List[str]: return [ @@ -507,6 +534,10 @@ def get_table_and_columns( """ pass + def get_data_element(self, data_element_name: str) -> Optional[DataElementTuple]: + """Override this to get data element by name""" + pass + @abstractclassmethod def get_metastore_params_template(self) -> AllFormField: """Override this to get the form field required for the metastore diff --git a/querybook/server/logic/data_element.py b/querybook/server/logic/data_element.py index b39a16c4a..e5fb0cdb8 100644 --- a/querybook/server/logic/data_element.py +++ b/querybook/server/logic/data_element.py @@ -1,3 +1,4 @@ +from typing import Union from app.db import with_session from const.data_element import ( DataElementAssociationProperty, @@ -61,23 +62,32 @@ def create_or_update_data_element( else: session.flush() + return data_element + @with_session -def create_data_element_association_by_name( - data_element_name: str, +def create_data_element_association( + metastore_id: int, + data_element_tuple: Union[DataElementTuple, str], column_id: int, association_type: DataElementAssociationType, property_name: DataElementAssociationProperty, primitive_type: str = None, session=None, ): - data_element = None - if data_element_name: - data_element = get_data_element_by_name(data_element_name, session=session) + if data_element_tuple is None: + return None + + if type(data_element_tuple) is str: + data_element = get_data_element_by_name(data_element_tuple, session=session) + else: + data_element = create_or_update_data_element( + metastore_id, data_element_tuple, session=session + ) if not data_element and not primitive_type: raise Exception( - f"Can not create DataElementAssociation: {data_element_name} is not a valid data element name and primitive type is empty" + f"Can not create DataElementAssociation: {data_element_tuple} is not a valid data element and primitive type is empty" ) return DataElementAssociation( @@ -91,6 +101,7 @@ def create_data_element_association_by_name( @with_session def create_column_data_element_association( + metastore_id: int, column_id: int, data_element_association: DataElementAssociationTuple, commit=True, @@ -102,8 +113,9 @@ def create_column_data_element_association( if data_element_association is not None: try: - value_association = create_data_element_association_by_name( - data_element_name=data_element_association.value, + value_association = create_data_element_association( + metastore_id=metastore_id, + data_element_tuple=data_element_association.value_data_element, column_id=column_id, association_type=data_element_association.type, property_name=DataElementAssociationProperty.VALUE, @@ -112,8 +124,9 @@ def create_column_data_element_association( ) if data_element_association.type == DataElementAssociationType.MAP: - key_association = create_data_element_association_by_name( - data_element_name=data_element_association.key, + key_association = create_data_element_association( + metastore_id=metastore_id, + data_element_tuple=data_element_association.key_data_element, column_id=column_id, association_type=data_element_association.type, property_name=DataElementAssociationProperty.KEY,