Add fixtures; minor changes

src/compub/models/base.py

@@ -10,6 +10,7 @@ from sqlalchemy.sql.schema import Column
 from sqlalchemy.sql.sqltypes import TIMESTAMP
 from sqlalchemy_utils.functions.orm import get_columns
 from sqlalchemy_utils.listeners import force_auto_coercion
+from sqlalchemy_utils.types.choice import Choice
 from sqlmodel.main import Field, SQLModel
 
 from compub.exceptions import NoDatabaseConfigured
@@ -18,42 +19,55 @@ from compub.settings import settings
 from sqlmodel import create_engine, Session
 
 
+DEFAULT_PK_TYPE = int
+
+
 class DB:
     def __init__(self):
         self._engine: AsyncEngine | Engine | None = None
-        self._session_maker: sessionmaker | None = None
+        self.session_maker: sessionmaker | None = None
 
     def start_engine(self) -> None:
         if settings.db_uri is None:
             raise NoDatabaseConfigured
         if settings.db_uri.scheme == 'sqlite':
-            self._engine = create_engine(settings.db_uri)
+            self._engine = create_engine(settings.db_uri, echo=True)
-            self._session_maker = sessionmaker(self.engine, expire_on_commit=False, class_=Session)
+            self.session_maker = sessionmaker(self.engine, expire_on_commit=False, class_=Session)
         else:
             self._engine = create_async_engine(settings.db_uri, future=True)
-            self._session_maker = sessionmaker(self.engine, expire_on_commit=False, class_=AsyncSession)
+            self.session_maker = sessionmaker(self.engine, expire_on_commit=False, class_=AsyncSession)
         force_auto_coercion()
 
     @property
-    def engine(self):
+    def engine(self) -> Engine:
         if self._engine is None:
             self.start_engine()
         assert isinstance(self._engine, (AsyncEngine, Engine))
         return self._engine
 
-    async def get_session(self) -> Session | AsyncSession:
+    async def get_session(self) -> AsyncSession:
-        if self._session_maker is None:
+        if self.session_maker is None:
             self.start_engine()
-        assert isinstance(self._session_maker, sessionmaker)
+        assert isinstance(self.session_maker.class_, AsyncSession)
-        session = self._session_maker()
+        session = self.session_maker()
         try:
             yield session
         finally:
             await session.close()
 
+    def get_session_blocking(self) -> Session:
+        if self.session_maker is None:
+            self.start_engine()
+        assert isinstance(self.session_maker.class_, Session)
+        session = self.session_maker()
+        try:
+            yield session
+        finally:
+            session.close()
+
 
 class AbstractBase(SQLModel):
-    id: Optional[int] = Field(default=None, primary_key=True)
+    id: Optional[DEFAULT_PK_TYPE] = Field(default=None, primary_key=True)
     date_created: Optional[datetime] = Field(
         default=None, sa_column=Column(TIMESTAMP(timezone=False), server_default=db_now())
     )
@@ -78,3 +92,7 @@ class AbstractBase(SQLModel):
     @staticmethod
     def get_non_repr_fields() -> list[str]:
         return ['id', 'date_created', 'date_updated']
+
+
+def get_choice_value(obj: Choice | str) -> str:
+    return obj.value if isinstance(obj, Choice) else obj

src/compub/models/companies.py

@@ -1,5 +1,5 @@
 from datetime import date
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 
 from slugify import slugify
 from sqlalchemy.engine.base import Connection
@@ -7,19 +7,24 @@ 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
+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
+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',
 ]
@@ -31,7 +36,7 @@ class LegalForm(AbstractBase, table=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)))
-    country: str = Field(sa_column=Column(CountryType))
+    country: Optional[str] = Field(sa_column=Column(CountryType))
 
     # Relationships
     subcategories: list['LegalFormSubcategory'] = Relationship(
@@ -44,13 +49,16 @@ class LegalForm(AbstractBase, table=True):
 
 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[int] = Field(
+    legal_form_id: Optional[PK] = Field(
         foreign_key='legal_form.id', default=None, nullable=False, index=True
     )
     legal_form: Optional[LegalForm] = Relationship(
@@ -66,18 +74,32 @@ class LegalFormSubcategory(AbstractBase, table=True):
 
 
 class CompanyIndustryLink(AbstractBase, table=True):
-    __tablename__ = 'company_industries'
+    __tablename__ = 'company_industry'
+    __table_args__ = (
+        Index('ux_company_industry', 'company_id', 'industry_id', unique=True),
+    )
 
     # Relationships
-    company_id: Optional[int] = Field(foreign_key='company.id', default=None, nullable=False, primary_key=True)
+    company_id: Optional[PK] = Field(foreign_key='company.id', default=None, nullable=False, primary_key=True)
-    industry_id: Optional[int] = Field(foreign_key='industry.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)
+    name: str = Field(max_length=255, nullable=False, index=True, sa_column_kwargs={'unique': True})
 
     # Relationships
     companies: list['Company'] = Relationship(
@@ -88,6 +110,24 @@ class Industry(AbstractBase, table=True):
         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'
 
@@ -96,28 +136,40 @@ class Company(AbstractBase, table=True):
     insolvent: bool = Field(default=False, nullable=False, index=True)
     founding_data: Optional[date]
     liquidation_date: Optional[date]
-    city: str = Field(max_length=255, index=True)  # TODO: Get rid of city; implement address properly
+    # TODO: Get rid of city; implement address properly
+    city: Optional[str] = Field(max_length=255, nullable=True, index=True)
 
     # Relationships
-    legal_form_id: Optional[int] = Field(
+    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[int] = Field(
+    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"<Company {self.id}>")
 
@@ -128,18 +180,22 @@ class Company(AbstractBase, table=True):
 
 class CompanyName(AbstractBase, table=True):
     __tablename__ = 'company_name'
-    __MAX_NAME_LENGTH__: int = 768
+    __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_NAME_LENGTH__, sa_column=Column(Unicode(__MAX_NAME_LENGTH__), nullable=False, index=True)
+        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[int] = Field(
+    company_id: Optional[PK] = Field(
         foreign_key='company.id', default=None, nullable=False, index=True
     )
     company: Optional[Company] = Relationship(

src/compub/models/courts.py

@@ -0,0 +1,234 @@
+from enum import Enum as EnumPy
+from typing import Optional
+
+from sqlalchemy.sql.schema import Column, Index
+from sqlalchemy.sql.sqltypes import Unicode, Enum as EnumSQL
+from sqlmodel.main import Field, Relationship
+
+from .base import AbstractBase, DEFAULT_PK_TYPE as PK
+from .companies import Company
+from .geography import StateProvince, Address
+
+
+__all__ = [
+    'CourtClass',
+    'CourtDepartment',
+    'Court',
+    'RegisterBranch',
+    'RegisterNumber',
+]
+
+
+# Amtsgericht -> Landgericht -> OLG -> BGH
+class CourtClass(AbstractBase, table=True):
+    __tablename__ = 'court_class'
+
+    # Fields
+    short: str = Field(
+        max_length=32,
+        nullable=False,
+        index=True
+    )
+    name: str = Field(
+        max_length=255,
+        sa_column=Column(Unicode(255))
+    )
+
+    # Relationships
+    courts: list['Court'] = Relationship(
+        back_populates='court_class',
+        sa_relationship_kwargs={'lazy': 'selectin'}
+    )
+
+    def __str__(self) -> str:
+        return str(self.name)
+
+
+class CourtDepartmentCourtLink(AbstractBase, table=True):
+    __tablename__ = 'court_department_court'
+    __table_args__ = (
+        Index('ux_court_department_court', 'court_department_id', 'court_id', unique=True),
+    )
+
+    # Relationships
+    court_department_id: Optional[PK] = Field(
+        foreign_key='court_department.id',
+        default=None,
+        nullable=False,
+        primary_key=True
+    )
+    court_id: Optional[PK] = Field(
+        foreign_key='court.id',
+        default=None,
+        nullable=False,
+        primary_key=True
+    )
+
+
+# Abteilung innerhalb eines Amtsgerichts (z. B. Registergericht, Insolvenzgericht etc.)
+class CourtDepartment(AbstractBase, table=True):
+    __tablename__ = 'court_department'
+
+    __REG_COURT_PK__ = 1
+    __INS_COURT_PK__ = 2
+
+    # Fields
+    name: str = Field(
+        max_length=255,
+        sa_column=Column(Unicode(255))
+    )
+
+    # Relationships
+    courts: list['Court'] = Relationship(
+        back_populates='court_departments',
+        link_model=CourtDepartmentCourtLink,
+        sa_relationship_kwargs={'lazy': 'selectin'}
+    )
+
+    def __str__(self) -> str:
+        return str(self.name)
+
+
+class Court(AbstractBase, table=True):
+    __tablename__ = 'court'
+
+    # Fields
+    name: str = Field(
+        max_length=255,
+        sa_column=Column(Unicode(255))
+    )
+
+    # Relationships
+    court_class_id: Optional[PK] = Field(
+        foreign_key='court_class.id',
+        default=None,
+        nullable=False,
+        index=True
+    )
+    court_class: Optional[CourtClass] = Relationship(
+        back_populates='courts',
+        sa_relationship_kwargs={'lazy': 'selectin'}
+    )
+
+    state_province_id: Optional[PK] = Field(
+        foreign_key='state_province.id',
+        default=None,
+        nullable=True,
+        index=True
+    )
+    state_province: Optional[StateProvince] = Relationship(
+        back_populates='courts',
+        sa_relationship_kwargs={'lazy': 'selectin'}
+    )
+
+    address_id: Optional[PK] = Field(
+        foreign_key='address.id',
+        default=None,
+        nullable=True,
+        index=True
+    )
+    address: Optional[Address] = Relationship(
+        back_populates='courts',
+        sa_relationship_kwargs={'lazy': 'selectin'}
+    )
+
+    parent_court_id: Optional[PK] = Field(
+        foreign_key='court.id',
+        default=None,
+        nullable=True,
+        index=True
+    )
+    parent_court: Optional['Court'] = Relationship(
+        back_populates='sub_courts',
+        sa_relationship_kwargs=dict(
+            lazy='selectin',
+            remote_side='Court.id'
+        )
+    )
+    sub_courts: list['Court'] = Relationship(
+        back_populates='parent_court',
+        sa_relationship_kwargs={'lazy': 'selectin'}
+    )
+
+    court_departments: list[CourtDepartment] = Relationship(
+        back_populates='courts',
+        link_model=CourtDepartmentCourtLink,
+        sa_relationship_kwargs={'lazy': 'selectin'}
+    )
+
+    register_numbers: list['RegisterNumber'] = Relationship(
+        back_populates='court', sa_relationship_kwargs={'lazy': 'selectin'}
+    )
+
+    def __str__(self) -> str:
+        return str(self.name)
+
+
+class RegisterBranch(str, EnumPy):
+    HRA = 'HRA'
+    HRB = 'HRB'
+    GnR = 'GnR'
+    PR = 'PR'
+    VR = 'VR'
+
+
+class RegisterNumber(AbstractBase, table=True):
+    __tablename__ = 'register_number'
+
+    # Fields
+    branch: Optional[RegisterBranch] = Field(
+        default=None,
+        sa_column=Column(
+            EnumSQL(RegisterBranch),
+            nullable=True,
+            index=True
+        )
+    )
+    number: str = Field(
+        max_length=255,
+        index=True
+    )
+    suffix: Optional[str] = Field(
+        max_length=4,
+        nullable=True
+    )
+    ureg_company_id: Optional[int] = Field(
+        nullable=True,
+        index=True
+    )
+
+    # Relationships
+    court_id: Optional[PK] = Field(
+        foreign_key='court.id',
+        default=None,
+        index=True
+    )
+    court: Optional[Court] = Relationship(
+        back_populates='register_numbers',
+        sa_relationship_kwargs=dict(
+            lazy='selectin'
+        )
+    )
+
+    company_id: Optional[PK] = Field(
+        foreign_key='company.id',
+        default=None,
+        index=True
+    )
+    company: Optional[Company] = Relationship(
+        back_populates='register_numbers',
+        sa_relationship_kwargs=dict(
+            lazy='selectin'
+        )
+    )
+
+    def __str__(self) -> str:
+        return f'{self.with_branch_code} ({self.company.current_name})'
+    
+    @property
+    def with_branch_code(self) -> str:
+        return f'{self.branch} {self.number}' if self.branch else str(self.number)
+    
+    @property
+    def verbose_id(self) -> str:
+        return f'{self.court} {self.with_branch_code}'

src/compub/models/geography.py

@@ -1,13 +1,17 @@
-from typing import Any, Optional
+from typing import Any, Optional, TYPE_CHECKING
 
 from pydantic import validator
-from sqlalchemy.sql.schema import Column
+from sqlalchemy.sql.schema import Column, Index
 from sqlalchemy.sql.sqltypes import Unicode
 from sqlalchemy_utils.primitives.country import Country
 from sqlalchemy_utils.types import CountryType
 from sqlmodel.main import Field, Relationship
 
-from .base import AbstractBase
+from .base import AbstractBase, DEFAULT_PK_TYPE as PK
+
+if TYPE_CHECKING:
+    from .companies import Company
+    from .courts import Court
 
 
 __all__ = [
@@ -20,6 +24,9 @@ __all__ = [
 
 class StateProvince(AbstractBase, table=True):
     __tablename__ = 'state_province'
+    __table_args__ = (
+        Index('ux_state_province_name_country', 'name', 'country', unique=True),
+    )
 
     # Fields
     country: str = Field(sa_column=Column(CountryType, nullable=False, index=True))
@@ -30,6 +37,10 @@ class StateProvince(AbstractBase, table=True):
         back_populates='state_province', sa_relationship_kwargs={'lazy': 'selectin'}
     )
 
+    courts: list['Court'] = Relationship(
+        back_populates='state_province', sa_relationship_kwargs={'lazy': 'selectin'}
+    )
+
     @validator('country', pre=True)
     def country_as_uppercase_string(cls, v: Any) -> str:
         if isinstance(v, Country):
@@ -41,13 +52,16 @@ class StateProvince(AbstractBase, table=True):
 
 class City(AbstractBase, table=True):
     __tablename__ = 'city'
+    __table_args__ = (
+        Index('ux_city_name_zip_state', 'name', 'zip_code', 'state_province_id', unique=True),
+    )
 
     # Fields
     zip_code: str = Field(max_length=5, nullable=False, index=True)
     name: str = Field(max_length=255, sa_column=Column(Unicode(255), nullable=False, index=True))
 
     # Relationships
-    state_province_id: Optional[int] = Field(
+    state_province_id: Optional[PK] = Field(
         foreign_key='state_province.id', default=None, nullable=False, index=True
     )
     state_province: Optional[StateProvince] = Relationship(
@@ -66,7 +80,7 @@ class Street(AbstractBase, table=True):
     name: str = Field(max_length=255, sa_column=Column(Unicode(255), nullable=False, index=True))
 
     # Relationships
-    city_id: Optional[int] = Field(
+    city_id: Optional[PK] = Field(
         foreign_key='city.id', default=None, nullable=False, index=True
     )
     city: Optional[City] = Relationship(
@@ -86,9 +100,17 @@ class Address(AbstractBase, table=True):
     supplement: str = Field(max_length=255)
 
     # Relationships
-    street_id: Optional[int] = Field(
+    street_id: Optional[PK] = Field(
         foreign_key='street.id', default=None, nullable=False, index=True
     )
     street: Optional[Street] = Relationship(
         back_populates='addresses', sa_relationship_kwargs={'lazy': 'selectin'}
     )
+
+    companies: list['Company'] = Relationship(
+        back_populates='address', sa_relationship_kwargs={'lazy': 'selectin'}
+    )
+
+    courts: list['Court'] = Relationship(
+        back_populates='address', sa_relationship_kwargs={'lazy': 'selectin'}
+    )

tests/fixtures/companies/company.json

@@ -0,0 +1,5 @@
+[
+  {
+
+  }
+]

tests/fixtures/companies/industry.json

@@ -0,0 +1,9 @@
+[
+  {
+    "id": 1,
+    "name": "Energy"
+  }, {
+    "id": 2,
+    "name": "Retail"
+  }
+]

tests/fixtures/companies/legal_form.json

@@ -0,0 +1,16 @@
+[
+  {
+    "id": 1,
+    "short": "KG",
+    "name": "Kommanditgesellschaft",
+    "country": "DE"
+  }, {
+    "id": 2,
+    "short": "AG",
+    "name": "Aktiengesellschaft",
+    "country": "DE"
+  }, {
+    "id": 123,
+    "short": "Ausl.HRB"
+  }
+]

tests/fixtures/companies/legal_form_subcategory.json

@@ -0,0 +1,17 @@
+[
+  {
+    "id": 111,
+    "short": "GmbH & Co. KG",
+    "name": "Gesellschaft mit beschränkter Haftung & Compagnie Kommanditgesellschaft",
+    "legal_form_id": 1
+  }, {
+    "id": 112,
+    "short": "AG & Co. KG",
+    "name": "Aktiengesellschaft & Compagnie Kommanditgesellschaft",
+    "legal_form_id": 1
+  }, {
+    "id": 69420,
+    "short": "LLC",
+    "legal_form_id": 123
+  }
+]

tests/fixtures/geography/address.json

@@ -0,0 +1,8 @@
+[
+  {
+    "id": 9999999,
+    "house_number": "123AB",
+    "supplement": "whatever",
+    "street_id": 1
+  }
+]

tests/fixtures/geography/city.json

@@ -0,0 +1,23 @@
+[
+  {
+    "id": 1122,
+    "zip_code": "10405",
+    "name": "Berlin",
+    "state_province_id": 100
+  }, {
+    "id": 9,
+    "zip_code": "33299",
+    "name": "Miami",
+    "state_province_id": 30
+  }, {
+    "id": 10,
+    "zip_code": "33694",
+    "name": "Tampa",
+    "state_province_id": 30
+  }, {
+    "id": 78787878,
+    "zip_code": "91000",
+    "name": "Luhansk",
+    "state_province_id": 4531354
+  }
+]

tests/fixtures/geography/state_province.json

@@ -0,0 +1,15 @@
+[
+  {
+    "id": 100,
+    "country": "DE",
+    "name": "Berlin"
+  }, {
+    "id": 30,
+    "country": "US",
+    "name": "Florida"
+  }, {
+    "id": 4531354,
+    "country": "UA",
+    "name": "Luhansk Oblast"
+  }
+]

tests/fixtures/geography/street.json

@@ -0,0 +1,7 @@
+[
+  {
+    "id": 1,
+    "name": "Some Street",
+    "city_id": 1122
+  }
+]