[0.2.0] Metadata Implementation

sdgx/data_models/metadata.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import json
-from enum import Enum
 from pathlib import Path
 from typing import Any, Dict, List
 
@@ -10,7 +9,7 @@
 
 from sdgx.data_loader import DataLoader
 from sdgx.data_models.inspectors.manager import InspectorManager
-from sdgx.exceptions import MetadataInitError
+from sdgx.exceptions import MetadaError, MetadataInitError
 from sdgx.utils import logger
 
 
@@ -22,13 +21,9 @@ class Metadata(BaseModel):
     For each column, there should be an instance of the Data Type object.
 
     Args:
-        primary_key(str): The primary key, a field used to uniquely identify each row in the table.
+        primary_keys(List[str]): The primary key, a field used to uniquely identify each row in the table.
         The primary key of each row must be unique and not empty.
 
-        composite_primary_key(bool): Whether to enable the composite primary key feature.
-
-        primary_key_list(bool): List of composite primary keys.
-
         column_list(list[str]): list of the comlumn name in the table, other columns lists are used to store column information.
     """
 
@@ -39,6 +34,7 @@ class Metadata(BaseModel):
     # variables related to columns
     # column_list is used to store all columns' name
     column_list: List[str] = []
+
     # other columns lists are used to store column information
     # here are 5 basic data types
     id_columns: List[str] = []
@@ -47,11 +43,9 @@ class Metadata(BaseModel):
     discrete_columns: List[str] = []
     datetime_columns: List[str] = []
 
-    # _column_dict = {}
-    _extend: Dict[str, Any] = {}
-
     # version info
     metadata_version: str = "1.0"
+    _extend: Dict[str, Any] = {}
 
     def get(self, key: str, default=None) -> Any:
         return getattr(self, key, getattr(self._extend, key, default))
@@ -76,7 +70,7 @@ def from_dataloader(
         cls,
         dataloader: DataLoader,
         max_chunk: int = 10,
-        primary_key: str = None,
+        primary_keys: List[str] = None,
         include_inspectors: list[str] | None = None,
         exclude_inspectors: list[str] | None = None,
         inspector_init_kwargs: dict[str, Any] | None = None,
@@ -107,11 +101,11 @@ def from_dataloader(
             if all(i.ready for i in inspectors) or i > max_chunk:
                 break
 
-        # If primary_key is not specified, use the first column.
-        if primary_key is None:
-            primary_key = dataloader.columns()[0]
+        # If primary_key is not specified, use the first column (in list).
+        if primary_keys is None:
+            primary_keys = [dataloader.columns()[0]]
 
-        metadata = Metadata(primary_key=primary_key, column_list=dataloader.columns())
+        metadata = Metadata(primary_keys=primary_keys, column_list=dataloader.columns())
         for inspector in inspectors:
             metadata.update(inspector.inspect())
 
@@ -131,7 +125,7 @@ def from_dataframe(
         for inspector in inspectors:
             inspector.fit(df)
 
-        metadata = Metadata(primary_key=df.columns[0], column_list=list(df.columns))
+        metadata = Metadata(primary_keys=[df.columns[0]], column_list=list(df.columns))
         for inspector in inspectors:
             metadata.update(inspector.inspect())
 
@@ -147,41 +141,79 @@ def load(cls, path: str | Path) -> "Metadata":
         attributes = json.load(path.open("r"))
         return Metadata().update(attributes)
 
+    def check_single_primary_key(self, input_key: str):
+        """Check whether a primary key in column_list and has ID data type.
+
+        Args:
+            input_key(str): the input primary_key str
+        """
+
+        if input_key not in self.column_list:
+            raise MetadaError(f"Primary Key {input_key} not Exist in columns.")
+        if input_key not in self.id_columns:
+            raise MetadaError(f"Primary Key {input_key} should has ID DataType.")
+
+    def get_all_data_type_columns(self):
+        """Get all column names from `self.xxx_columns`.
+
+        All Lists with the suffix _columns in model fields and extend fields need to be collected.
+        All defined column names will be counted.
+
+        Returns:
+            all_dtype_cols(set): set of all column names.
+        """
+        all_dtype_cols = set()
+
+        # search the model fields and extend fields
+        for each_key in list(self.model_fields.keys()) + list(self._extend.keys()):
+            if each_key.endswith("_columns"):
+                column_names = self.get(each_key)
+                all_dtype_cols = all_dtype_cols.union(set(column_names))
+
+        return all_dtype_cols
+
     def check(self):
         """Checks column info.
 
         When passing as input to the next module, perform necessary checks, including:
-            -Is the primary key correctly defined.
-            -Is there any missing definition of the column.
-            -Are there any  unknown columns that have been incorrectly updated.
+            -Is the primary key correctly defined(in column list) and has ID data type.
+            -Is there any missing definition of each column in table.
+            -Are there any unknown columns that have been incorrectly updated.
         """
-        # Not implemented yet
+        # check primary key in column_list and has ID data type
+        for each_key in self.primary_keys:
+            self.check_single_primary_key(each_key)
 
-        pass
+        all_dtype_columns = self.get_all_data_type_columns()
 
-    def update_primary_key(self, primary_key: str | list[str], composite_primary_key: bool = False):
+        # check missing columns
+        for each_column in self.column_list:
+            if each_column not in all_dtype_columns:
+                raise MetadaError(f"Undefined data type for column {each_column}.")
+
+        # check unfamiliar columns in dtypes
+        for each_dtype_column in all_dtype_columns:
+            if each_dtype_column not in self.column_list:
+                raise MetadaError(f"Found undefined column: {each_dtype_column}.")
+
+        logger.info("Metadata check finished.")
+
+    def update_primary_key(self, primary_keys: List[str]):
         """Update the primary key of the table
 
         When update the primary key, the original primary key will be erased.
 
         Args:
-            primary_key(str | list[str]): the primary key or key list.
-
-            composite_primary_key(bool): whether this table use composite primary key.
+            primary_keys(List[str]): the primary keys of this table.
         """
 
-        if composite_primary_key is False and not isinstance(primary_key, str):
-            raise ValueError("Primary key should be a string")
+        if not isinstance(primary_keys, List):
+            raise MetadaError("Primary key should be a list.")
 
-        if composite_primary_key is True and len(primary_key) == 0:
-            raise ValueError("Composite primary key list shoud NOT be empty.")
+        for each_key in primary_keys:
+            if each_key not in self.column_list:
+                raise MetadaError("Primary key not exist in table columns.")
 
-        if composite_primary_key is True:
-            self._composite_primary_key = True
-            self.primary_key = None
-            self.primary_key_list = primary_key
-        else:
-            self._composite_primary_key = False
-            self.primary_key = primary_key
+        self.primary_keys = primary_keys
 
-        logger.info(f"Primary Key updated: {primary_key}.")
+        logger.info(f"Primary Key updated: {primary_keys}.")

sdgx/data_models/multi_table_combiner.py

@@ -4,6 +4,8 @@
 
 from sdgx.data_models.metadata import Metadata
 from sdgx.data_models.relationship import Relationship
+from sdgx.exceptions import MultiTableCombinerError
+from sdgx.utils import logger
 
 
 class MultiTableCombiner(BaseModel):
@@ -31,7 +33,7 @@ def check(self):
         relationship_cnt = len(self.relationships)
         metadata_cnt = len(self.metadata_dict.keys())
         if metadata_cnt != relationship_cnt + 1:
-            raise ValueError("Number of tables should corresponds to relationships.")
+            raise MultiTableCombinerError("Number of tables should corresponds to relationships.")
 
         # table name check
         table_names_from_relationships = set()
@@ -40,15 +42,19 @@ def check(self):
         table_names = list(self.metadata_dict.keys())
         for each_r in self.relationships:
             if each_r.parent_table not in table_names:
-                raise ValueError(f"Metadata of parent table {each_r.parent_table} is missing.")
+                raise MultiTableCombinerError(
+                    f"Metadata of parent table {each_r.parent_table} is missing."
+                )
             if each_r.child_table not in table_names:
-                raise ValueError(f"Metadata of child table {each_r.child_table} is missing.")
+                raise MultiTableCombinerError(
+                    f"Metadata of child table {each_r.child_table} is missing."
+                )
             table_names_from_relationships.add(each_r.parent_table)
             table_names_from_relationships.add(each_r.child_table)
 
         # each table in metadata must in a relationship
         for each_t in table_names:
             if each_t not in table_names_from_relationships:
-                raise ValueError(f"Table {each_t} has not relationship.")
+                raise MultiTableCombinerError(f"Table {each_t} has not relationship.")
 
-        return True
+        logger.info("MultiTableCombiner check finished.")

sdgx/data_models/relationship.py

@@ -2,6 +2,8 @@
 
 from pydantic import BaseModel
 
+from sdgx.exceptions import RelationshipError
+
 
 class Relationship(BaseModel):
     """Relationship between tables
@@ -18,15 +20,10 @@ class Relationship(BaseModel):
     parent_table: str
     child_table: str
 
-    # foreign keys
-    child_table_foreign_key: str = "foreign key undefined"
-
-    # for composite keys
-    composite_foreign_key: bool = False
-    child_table_composite_foreign_key: List[str] = []
+    foreign_keys: List[str] = []
 
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
         if self.parent_table == self.child_table:
-            raise ValueError("child table and parent table cannot be the same")
+            raise RelationshipError("child table and parent table cannot be the same")

sdgx/exceptions.py

@@ -108,3 +108,24 @@ class CannotExportError(SdgxError):
 
     EXIT_CODE = 107
     ERROR_CODE = 8001
+
+
+class DataModelError(SdgxError):
+    """
+    Exception to indicate that exception in all data models.
+    """
+
+    EXIT_CODE = 108
+    ERROR_CODE = 9001
+
+
+class MetadaError(DataModelError):
+    ERROR_CODE = 9002
+
+
+class RelationshipError(DataModelError):
+    ERROR_CODE = 9003
+
+
+class MultiTableCombinerError(DataModelError):
+    ERROR_CODE = 9004