from datetime import date from typing import Optional, TYPE_CHECKING from slugify import slugify from sqlalchemy.engine.base import Connection from sqlalchemy.event.api import listens_for from sqlalchemy.orm.mapper import Mapper from sqlalchemy.sql.expression import select from sqlalchemy.sql.functions import count from sqlalchemy.sql.schema import Column, Index from sqlalchemy.sql.sqltypes import Unicode from sqlalchemy_utils.types import CountryType from sqlmodel.main import Field, Relationship from compub.utils import multi_max from .base import AbstractBase, DEFAULT_PK_TYPE as PK from .geography import Address if TYPE_CHECKING: from .courts import RegisterNumber __all__ = [ 'LegalForm', 'LegalFormSubcategory', 'Industry', 'Executive', 'Company', 'CompanyName', ] class LegalForm(AbstractBase, table=True): __tablename__ = 'legal_form' # Fields short: str = Field(max_length=32, nullable=False, index=True) name: Optional[str] = Field(default=None, max_length=255, sa_column=Column(Unicode(255))) country: str = Field(sa_column=Column(CountryType)) # Relationships subcategories: list['LegalFormSubcategory'] = Relationship( back_populates='legal_form', sa_relationship_kwargs={'lazy': 'selectin'} ) def __str__(self) -> str: return str(self.short) class LegalFormSubcategory(AbstractBase, table=True): __tablename__ = 'legal_form_subcategory' __table_args__ = ( Index('ux_legal_form_subcategory', 'short', 'legal_form_id', unique=True), ) # Fields short: str = Field(max_length=32, nullable=False, index=True) name: Optional[str] = Field(default=None, max_length=255, sa_column=Column(Unicode(255))) # Relationships legal_form_id: Optional[PK] = Field( foreign_key='legal_form.id', default=None, nullable=False, index=True ) legal_form: Optional[LegalForm] = Relationship( back_populates='subcategories', sa_relationship_kwargs={'lazy': 'selectin'} ) companies: list['Company'] = Relationship( back_populates='legal_form', sa_relationship_kwargs={'lazy': 'selectin'} ) def __str__(self) -> str: return str(self.short) class CompanyIndustryLink(AbstractBase, table=True): __tablename__ = 'company_industry' __table_args__ = ( Index('ux_company_industry', 'company_id', 'industry_id', unique=True), ) # Relationships company_id: Optional[PK] = Field(foreign_key='company.id', default=None, nullable=False, primary_key=True) industry_id: Optional[PK] = Field(foreign_key='industry.id', default=None, nullable=False, primary_key=True) class CompanyExecutiveLink(AbstractBase, table=True): __tablename__ = 'company_executive' __table_args__ = ( Index('ux_company_executive', 'company_id', 'executive_id', unique=True), ) # Relationships company_id: Optional[PK] = Field(foreign_key='company.id', default=None, nullable=False, primary_key=True) executive_id: Optional[PK] = Field(foreign_key='executive.id', default=None, nullable=False, primary_key=True) class Industry(AbstractBase, table=True): __tablename__ = 'industry' # Fields name: str = Field(max_length=255, nullable=False, index=True, sa_column_kwargs={'unique': True}) # Relationships companies: list['Company'] = Relationship( back_populates='industries', link_model=CompanyIndustryLink, sa_relationship_kwargs={'lazy': 'selectin'} ) def __str__(self) -> str: return str(self.name) class Executive(AbstractBase, table=True): __tablename__ = 'executive' __MAX_LENGTH_NAME__: int = 255 # Fields name: str = Field( max_length=__MAX_LENGTH_NAME__, sa_column=Column(Unicode(__MAX_LENGTH_NAME__), nullable=False, index=True) ) # Relationships companies: list['Company'] = Relationship( back_populates='executives', link_model=CompanyExecutiveLink, sa_relationship_kwargs={'lazy': 'selectin'} ) def __str__(self) -> str: return str(self.name) class Company(AbstractBase, table=True): __tablename__ = 'company' # Fields visible: bool = Field(default=True, nullable=False, index=True) insolvent: bool = Field(default=False, nullable=False, index=True) founding_data: Optional[date] liquidation_date: Optional[date] # TODO: Get rid of city; implement address properly city: Optional[str] = Field(max_length=255, nullable=True, index=True) # Relationships legal_form_id: Optional[PK] = Field( foreign_key='legal_form_subcategory.id', default=None, index=True ) legal_form: Optional[LegalFormSubcategory] = Relationship( back_populates='companies', sa_relationship_kwargs={'lazy': 'selectin'} ) address_id: Optional[PK] = Field( foreign_key='address.id', default=None, index=True ) address: Optional[Address] = Relationship( back_populates='companies', sa_relationship_kwargs={'lazy': 'selectin'} ) industries: list[Industry] = Relationship( back_populates='companies', link_model=CompanyIndustryLink, sa_relationship_kwargs={'lazy': 'selectin'} ) executives: list[Executive] = Relationship( back_populates='companies', link_model=CompanyExecutiveLink, sa_relationship_kwargs={'lazy': 'selectin'} ) names: list['CompanyName'] = Relationship( back_populates='company', sa_relationship_kwargs={'lazy': 'selectin'} ) register_numbers: list['RegisterNumber'] = Relationship( back_populates='company', sa_relationship_kwargs={'lazy': 'selectin'} ) def __str__(self) -> str: return str(self.current_name or f"") @property def current_name(self) -> Optional['CompanyName']: return multi_max(list(self.names), CompanyName.get_reg_date, 'date_updated', default=None) class CompanyName(AbstractBase, table=True): __tablename__ = 'company_name' __table_args__ = ( Index('ux_company_name_company_id', 'name', 'company_id', unique=True), ) __MAX_LENGTH_NAME__: int = 768 __MAX_SLUG_LENGTH__: int = 255 # Fields name: str = Field( max_length=__MAX_LENGTH_NAME__, sa_column=Column(Unicode(__MAX_LENGTH_NAME__), nullable=False, index=True) ) date_registered: Optional[date] slug: Optional[str] = Field(default=None, max_length=__MAX_SLUG_LENGTH__, index=True) # Relationships company_id: Optional[PK] = Field( foreign_key='company.id', default=None, nullable=False, index=True ) company: Optional[Company] = Relationship( back_populates='names', sa_relationship_kwargs={'lazy': 'selectin'} ) def __str__(self) -> str: return str(self.name) def get_reg_date(self) -> date: return date(1, 1, 1) if self.date_registered is None else self.date_registered @property def max_slug_length(self) -> int: return self.__MAX_SLUG_LENGTH__ @listens_for(CompanyName, 'before_insert') def generate_company_name_slug(_mapper: Mapper, connection: Connection, target: CompanyName) -> None: if target.slug: return slug = slugify(target.name)[:(target.max_slug_length - 2)] statement = select(count()).select_from(CompanyName).where(CompanyName.slug.startswith(slug)) num = connection.execute(statement).scalar() if num == 0: target.slug = slug else: target.slug = f'{slug}-{str(num + 1)}'