Anton_wireframe/app/db/models.py

import datetime
import enum

from sqlalchemy import Column, DateTime, ForeignKey, Integer, String, Table, Text
from sqlalchemy.orm import relationship
from sqlalchemy.types import Enum

from db.db import Base


class InvestmentStage(enum.Enum):
    SEED = "SEED"
    SERIES_A = "SERIES_A"
    SERIES_B = "SERIES_B"
    SERIES_C = "SERIES_C"
    GROWTH = "GROWTH"
    LATE_STAGE = "LATE_STAGE"

# Association table for many-to-many relationship between investors and companies
investor_company_association = Table(
    "investor_companies",
    Base.metadata,
    Column("investor_id", Integer, ForeignKey("investors.id")),
    Column("company_id", Integer, ForeignKey("companies.id")),
)


# Association table for investor-sector many-to-many
investor_sector_association = Table(
    "investor_sectors",
    Base.metadata,
    Column("investor_id", Integer, ForeignKey("investors.id")),
    Column("sector_id", Integer, ForeignKey("sectors.id")),
)


class InvestorTable(Base):
    __tablename__ = "investors"

    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, nullable=False)
    description = Column(Text, nullable=True)
    aum = Column(Integer, nullable=False)  # Assets Under Management
    check_size_lower = Column(Integer, nullable=False)  # Lower bound
    check_size_upper = Column(Integer, nullable=False)  # Upper bound
    geographic_focus = Column(String, nullable=False)
    stage_focus = Column(Enum(InvestmentStage), nullable=False)
    number_of_investments = Column(Integer, default=0)
    created_at = Column(DateTime, default=datetime.datetime.now(datetime.UTC))
    updated_at = Column(
        DateTime,
        default=datetime.datetime.now(datetime.UTC),
        onupdate=datetime.datetime.now(datetime.UTC),
    )

    # Relationship to portfolio companies
    portfolio_companies = relationship(
        "CompanyTable",
        secondary=investor_company_association,
        back_populates="investors",
    )
    team_members = relationship("InvestorTeamMember", back_populates="investor")
    sectors = relationship(
        "SectorTable",
        secondary=investor_sector_association,
        back_populates="investors",
    )


class CompanyTable(Base):
    __tablename__ = "companies"

    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, nullable=False)
    industry = Column(String, nullable=False)
    location = Column(String, nullable=False)
    founded_year = Column(Integer, nullable=True)
    website = Column(String, nullable=True)
    created_at = Column(DateTime, default=datetime.datetime.now(datetime.UTC))
    updated_at = Column(
        DateTime,
        default=datetime.datetime.now(datetime.UTC),
        onupdate=datetime.datetime.now(datetime.UTC),
    )

    # Relationship back to investors
    investors = relationship(
        "InvestorTable",
        secondary=investor_company_association,
        back_populates="portfolio_companies",
    )


class SectorTable(Base):
    __tablename__ = "sectors"

    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, nullable=False)

    # Add relationship back to investors
    investors = relationship(
        "InvestorTable",
        secondary=investor_sector_association,
        back_populates="sectors",
    )


class InvestorTeamMember(Base):
    __tablename__ = "investor_team"
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String, nullable=False)
    role = Column(String, nullable=False)
    email = Column(String, nullable=False)

    investor_id = Column(Integer, ForeignKey("investors.id"))
    investor = relationship("InvestorTable", back_populates="team_members")