diff --git a/src/app/db/models/custom_types.py b/src/app/db/models/custom_types.py index c2e1b249..fc64c3da 100644 --- a/src/app/db/models/custom_types.py +++ b/src/app/db/models/custom_types.py @@ -6,7 +6,7 @@ from sqlalchemy.types import TypeDecorator, String from sqlalchemy.dialects.postgresql import JSONB -from app.lib.schema import Location, Funding, WorkExperience, SocialActivity, OpportunityStage +from app.lib.schema import Location, FundingRound, Funding, WorkExperience, SocialActivity, OpportunityStage class JSONBType(TypeDecorator): @@ -16,7 +16,7 @@ def process_bind_param(self, value, dialect): """Convert Python object to JSON format before storing it in the database.""" if isinstance(value, dict): return value - elif hasattr(value, 'to_dict'): + elif hasattr(value, "to_dict"): return value.to_dict() return value @@ -39,7 +39,8 @@ class FundingType(JSONBType): def process_result_value(self, value, dialect): """Convert JSON format to Python object when reading from the database.""" if value and isinstance(value, dict): - return Funding.from_dict(value) + obj = Funding.from_dict(value) + obj.round_name = FundingRound(obj.round_name) if obj.round_name else FundingRound.SERIES_UNKNOWN return None @@ -50,6 +51,7 @@ def process_result_value(self, value, dialect): return WorkExperience.from_dict(value) return None + class SocialActivityType(JSONBType): def process_result_value(self, value, dialect): """Convert JSON format to Python object when reading from the database.""" diff --git a/src/app/lib/schema.py b/src/app/lib/schema.py index 5172d7c8..490ce3bf 100644 --- a/src/app/lib/schema.py +++ b/src/app/lib/schema.py @@ -26,6 +26,7 @@ class Message(CamelizedBaseStruct): class Location(CamelizedBaseStruct): """A Location.""" + city: str | None = None region: str | None = None country: str | None = None @@ -33,15 +34,32 @@ class Location(CamelizedBaseStruct): class Investor(CamelizedBaseStruct): """An investor.""" + name: str type: str | None = None url: str | None = None linkedin_profile_url: str | None = None +class FundingRound(enum.Enum): + """Funding round.""" + + GRANT = "Grant" + PRE_SEED = "Pre-Seed" + SEED = "Seed" + SERIES_A = "Series A" + SERIES_B = "Series B" + SERIES_C = "Series C" + SERIES_D = "Series D" + SERIES_E = "Series E" + SERIES_UNKNOWN = "Series Unknown" + PUBLIC = "Public" + + class Funding(CamelizedBaseStruct): """Funding data.""" - round_name: str = "Series Unknown" + + round_name: FundingRound = FundingRound.SERIES_UNKNOWN money_raised: int | None = None announced_date: date | None = None investors: list[Investor] = [] @@ -49,6 +67,7 @@ class Funding(CamelizedBaseStruct): class WorkExperience(CamelizedBaseStruct): """Work experience data.""" + starts_at: date title: str company_name: str @@ -61,6 +80,7 @@ class WorkExperience(CamelizedBaseStruct): class SocialActivity(CamelizedBaseStruct): """Social activity data.""" + title: str link: str | None = None status: str | None = None @@ -68,12 +88,13 @@ class SocialActivity(CamelizedBaseStruct): class OpportunityStage(enum.Enum): """Opportunity stages.""" - IDENTIFIED = "identified" - QUALIFIED = "qualified" - CONTACTED = "contacted" - ENGAGED = "engaged" - PROPOSED = "proposed" - NEGOTIATED = "negotiated" - DEFERRED = "deferred" - SUSPENDED = "suspended" - CUSTOMER = "customer" + + IDENTIFIED = "Identified" + QUALIFIED = "Qualified" + CONTACTED = "Contacted" + ENGAGED = "Engaged" + PROPOSED = "Proposed" + NEGOTIATED = "Negotiated" + DEFERRED = "Deferred" + SUSPENDED = "Suspended" + CUSTOMER = "Customer"