diff --git a/definitions/models.py b/definitions/models.py index fd3654c..ba36470 100644 --- a/definitions/models.py +++ b/definitions/models.py @@ -1,47 +1,54 @@ from flask_sqlalchemy import SQLAlchemy -from sqlalchemy import Enum, Integer, Float, String, Date, Column +from sqlalchemy import Enum, Integer, Float, String, Date, Column, Boolean from config import item_attributes, sql_conf from definitions.attribute import textAttribute, intAttribute, floatAttribute, dateAttribute, selectAttribute # Initialize SQLAlchemy db = SQLAlchemy() +# Mapping of attribute types to SQLAlchemy column types +COLUMN_TYPE_MAPPING = { + textAttribute: lambda attrib: String(attrib.max_length), + dateAttribute: lambda attrib: Date, + selectAttribute: lambda attrib: Enum(*attrib.options), + intAttribute: lambda attrib: Integer, + floatAttribute: lambda attrib: Float, +} + +def _get_column_type(attrib): + """Helper function to get the SQLAlchemy column type for an attribute.""" + if type(attrib) not in COLUMN_TYPE_MAPPING: + raise ValueError(f"Unsupported attribute type: {type(attrib)} for attribute '{attrib.attrib_name}'") + return COLUMN_TYPE_MAPPING[type(attrib)](attrib) + +def _get_column_properties(attrib): + """Helper function to get the properties for a column.""" + return { + 'primary_key': attrib.primary, + 'unique': attrib.unique, + 'nullable': not attrib.required, + 'default': attrib.default_val + } + def create_asset_model(): """ Dynamically creates the Asset model based on the configuration in item_attributes. """ # Define the table name and basic methods attrs = { - '__tablename__': sql_conf.SQL_TABLE, # Table name from config - '__init__': lambda self, **kwargs: self.__dict__.update(kwargs), # Constructor - '__repr__': lambda self: f"" # Representation + '__tablename__': sql_conf.SQL_TABLE, + '__init__': lambda self, **kwargs: self.__dict__.update(kwargs), + '__repr__': lambda self: f"" } # Define columns based on item_attributes for attrib in item_attributes: - # Determine the column type - if isinstance(attrib, textAttribute): - column_type = String(attrib.max_length) - elif isinstance(attrib, dateAttribute): - column_type = Date - elif isinstance(attrib, selectAttribute): - column_type = Enum(*attrib.options) - elif isinstance(attrib, (intAttribute, floatAttribute)): - column_type = Integer if isinstance(attrib, intAttribute) else Float - else: - raise ValueError(f"Unsupported attribute type: {type(attrib)}") - - # Define the column with additional properties - attrs[attrib.attrib_name] = Column( - column_type, - primary_key=attrib.primary, - unique=attrib.unique, - nullable=not attrib.required, - default=attrib.default_val - ) + column_type = _get_column_type(attrib) + column_properties = _get_column_properties(attrib) + attrs[attrib.attrib_name] = Column(column_type, **column_properties) # Create the Item class dynamically - Item = type('Asset', (db.Model,), attrs) + Item = type('Item', (db.Model,), attrs) return Item # Create the Item model