Add Company.org_size and JobPost.tools

src/app/db/migrations/versions/2024-08-15_add_company_org_size_and_jobpost_tools_0d39d07c03b2.py

@@ -0,0 +1,84 @@
+# type: ignore
+"""Add Company.org_size and JobPost.tools
+
+Revision ID: 0d39d07c03b2
+Revises: 46acf6e88bda
+Create Date: 2024-08-15 11:23:48.004899+00:00
+
+"""
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING
+
+import sqlalchemy as sa
+from alembic import op
+from advanced_alchemy.types import EncryptedString, EncryptedText, GUID, ORA_JSONB, DateTimeUTC
+from sqlalchemy import Text  # noqa: F401
+
+from app.db.models.custom_types import OrgSizeType, ToolType
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+__all__ = ["downgrade", "upgrade", "schema_upgrades", "schema_downgrades", "data_upgrades", "data_downgrades"]
+
+sa.GUID = GUID
+sa.DateTimeUTC = DateTimeUTC
+sa.ORA_JSONB = ORA_JSONB
+sa.EncryptedString = EncryptedString
+sa.EncryptedText = EncryptedText
+
+# revision identifiers, used by Alembic.
+revision = "0d39d07c03b2"
+down_revision = "46acf6e88bda"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=UserWarning)
+        with op.get_context().autocommit_block():
+            schema_upgrades()
+            data_upgrades()
+
+
+def downgrade() -> None:
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=UserWarning)
+        with op.get_context().autocommit_block():
+            data_downgrades()
+            schema_downgrades()
+
+
+def schema_upgrades() -> None:
+    """schema upgrade migrations go here."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("company", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("org_size", OrgSizeType(), nullable=True))
+
+    with op.batch_alter_table("job_post", schema=None) as batch_op:
+        batch_op.add_column(sa.Column("tools", ToolType(), nullable=True))
+
+    # ### end Alembic commands ###
+
+
+def schema_downgrades() -> None:
+    """schema downgrade migrations go here."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table("job_post", schema=None) as batch_op:
+        batch_op.drop_column("tools")
+
+    with op.batch_alter_table("company", schema=None) as batch_op:
+        batch_op.drop_column("org_size")
+
+    # ### end Alembic commands ###
+
+
+def data_upgrades() -> None:
+    """Add any optional data upgrade migrations here!"""
+
+
+def data_downgrades() -> None:
+    """Add any optional data downgrade migrations here!"""

src/app/db/models/company.py

@@ -7,8 +7,8 @@
 from sqlalchemy import String, ForeignKey
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
-from app.lib.schema import Location, Funding
-from .custom_types import LocationType, FundingType
+from app.lib.schema import Location, Funding, OrgSize
+from .custom_types import LocationType, FundingType, OrgSizeType
 
 
 class CompanyOrg(UUIDAuditBase):
@@ -37,6 +37,7 @@ class Company(UUIDAuditBase, SlugKey):
     linkedin_profile_url: Mapped[str | None] = mapped_column(String(length=2083), nullable=True, default=None)
     hq_location: Mapped[Location | None] = mapped_column(LocationType, nullable=True, default=None)
     last_funding: Mapped[Funding | None] = mapped_column(FundingType, nullable=True, default=None)
+    org_size: Mapped[OrgSize | None] = mapped_column(OrgSizeType, nullable=True, default=None)
     # -----------
     # ORM Relationships
     # ------------

src/app/db/models/custom_types.py

@@ -6,7 +6,17 @@
 from sqlalchemy.types import TypeDecorator, String
 from sqlalchemy.dialects.postgresql import JSONB
 
-from app.lib.schema import Location, FundingRound, Funding, WorkExperience, SocialActivity, OpportunityStage
+from app.lib.schema import (
+    Location,
+    FundingRound,
+    Funding,
+    WorkExperience,
+    SocialActivity,
+    OpportunityStage,
+    OrgSize,
+    Tool,
+    Scale,
+)
 
 
 class JSONBType(TypeDecorator):
@@ -73,3 +83,21 @@ def process_result_value(self, value, dialect):
         if value is not None:
             return OpportunityStage(value)
         return value
+
+
+class OrgSizeType(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 OrgSize.from_dict(value)
+        return None
+
+
+class ToolType(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):
+            obj = Tool.from_dict(value)
+            obj.certainty = Scale(obj.certainty) if obj.certainty else Scale.LOW
+            return obj
+        return None

src/app/db/models/job_post.py

@@ -8,8 +8,8 @@
 from sqlalchemy.dialects.postgresql import JSONB
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
-from app.lib.schema import Location, Funding
-from .custom_types import LocationType, FundingType
+from app.lib.schema import Location, Funding, Tool
+from .custom_types import LocationType, FundingType, ToolType
 from .company import Company
 
 
@@ -28,6 +28,7 @@ class JobPost(UUIDAuditBase):
     url: Mapped[str | None] = mapped_column(String(length=2083), nullable=True, default=None)
     apply_url: Mapped[str | None] = mapped_column(String(length=2083), nullable=True, default=None)
     external_id: Mapped[str | None] = mapped_column(nullable=True, default=None)
+    tools: Mapped[list[Tool] | None] = mapped_column(ToolType, nullable=True, default=None)
     company_id: Mapped[UUID] = mapped_column(ForeignKey("company.id"), nullable=True)
     # -----------
     # ORM Relationships

src/app/domain/companies/schemas.py

@@ -6,7 +6,7 @@
 import msgspec
 
 from app.db.models.company import Company
-from app.lib.schema import CamelizedBaseStruct, Location, Funding
+from app.lib.schema import CamelizedBaseStruct, Location, Funding, OrgSize
 
 
 class Company(CamelizedBaseStruct):
@@ -27,6 +27,7 @@ class Company(CamelizedBaseStruct):
     linkedin_profile_url: str | None = None
     hq_location: Location | None = None
     last_funding: Funding | None = None
+    org_size: OrgSize | None = None
 
 
 class CompanyCreate(CamelizedBaseStruct):

src/app/domain/jobs/schemas.py

@@ -6,7 +6,7 @@
 import msgspec
 
 from app.db.models.job_post import JobPost
-from app.lib.schema import CamelizedBaseStruct, Location
+from app.lib.schema import CamelizedBaseStruct, Location, Tool
 from app.domain.companies.schemas import Company
 
 
@@ -26,6 +26,7 @@ class JobPost(CamelizedBaseStruct):
     apply_url: str | None = None
     total_applicants: int | None = None
     external_id: str | None = None
+    tools: list[Tool] | None = None
     company: Company | None = None
 
 

src/app/lib/schema.py

@@ -98,3 +98,25 @@ class OpportunityStage(enum.Enum):
     DEFERRED = "Deferred"
     SUSPENDED = "Suspended"
     CUSTOMER = "Customer"
+
+
+class OrgSize(CamelizedBaseStruct):
+    """Org size data."""
+
+    engineering: int | None = None
+
+
+class Scale(CamelizedBaseStruct):
+    """Scale data."""
+
+    LOW = "Low"
+    MEDIUM = "Medium"
+    HIGH = "High"
+
+
+class Tool(CamelizedBaseStruct):
+    """Tool data."""
+
+    name: str
+    # Shouldn't be part of the tool but makes things easier for now
+    certainty: Scale = Scale.LOW