diff --git a/alembic/versions/2025_09_15_1137-d5f92e6fedf4_add_location_tables.py b/alembic/versions/2025_09_15_1137-d5f92e6fedf4_add_location_tables.py new file mode 100644 index 00000000..be2c22e9 --- /dev/null +++ b/alembic/versions/2025_09_15_1137-d5f92e6fedf4_add_location_tables.py @@ -0,0 +1,161 @@ +"""Add Location tables + +Revision ID: d5f92e6fedf4 +Revises: e7189dc92a83 +Create Date: 2025-09-15 11:37:58.183674 + +""" +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision: str = 'd5f92e6fedf4' +down_revision: Union[str, None] = 'e7189dc92a83' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + +US_STATES_TABLE_NAME = 'us_states' +COUNTIES_TABLE_NAME = 'counties' +LOCALITIES_TABLE_NAME = 'localities' +LOCATIONS_TABLE_NAME = 'locations' +LINK_AGENCIES_LOCATIONS_TABLE_NAME = 'link_agencies_locations' + +def upgrade() -> None: + _create_location_type() + _create_us_states_table() + _create_counties_table() + _create_localities_table() + _create_locations_table() + _create_link_agencies_locations_table() + +def downgrade() -> None: + _remove_link_agencies_locations_table() + _remove_locations_table() + _remove_localities_table() + _remove_counties_table() + _remove_us_states_table() + _remove_location_type() + +def _create_location_type(): + op.execute(""" + create type location_type as enum ('National', 'State', 'County', 'Locality') + """) + +def _remove_location_type(): + op.execute(""" + drop type location_type + """) + +def _create_us_states_table(): + op.execute(""" + create table if not exists public.us_states + ( + state_iso text not null + constraint unique_state_iso + unique, + state_name text, + id bigint generated always as identity + primary key + ) + """) + +def _create_counties_table(): + op.execute(""" + create table if not exists public.counties + ( + fips varchar not null + constraint unique_fips + unique, + name text, + lat double precision, + lng double precision, + population bigint, + agencies text, + id bigint generated always as identity + primary key, + state_id integer + references public.us_states, + unique (fips, state_id), + constraint unique_county_name_and_state + unique (name, state_id) + ) + """) + +def _create_localities_table(): + op.execute(""" + create table if not exists public.localities + ( + id bigint generated always as identity + primary key, + name varchar(255) not null + constraint localities_name_check + check ((name)::text !~~ '%,%'::text), + county_id integer not null + references public.counties, + unique (name, county_id) + ) + + """) + +def _create_locations_table(): + op.execute(""" + create table if not exists public.locations + ( + id bigint generated always as identity + primary key, + type location_type not null, + state_id bigint + references public.us_states + on delete cascade, + county_id bigint + references public.counties + on delete cascade, + locality_id bigint + references public.localities + on delete cascade, + lat double precision, + lng double precision, + unique (id, type, state_id, county_id, locality_id), + constraint locations_check + check (((type = 'National'::location_type) AND (state_id IS NULL) AND (county_id IS NULL) AND + (locality_id IS NULL)) OR + ((type = 'State'::location_type) AND (county_id IS NULL) AND (locality_id IS NULL)) OR + ((type = 'County'::location_type) AND (county_id IS NOT NULL) AND (locality_id IS NULL)) OR + ((type = 'Locality'::location_type) AND (county_id IS NOT NULL) AND (locality_id IS NOT NULL))) + ) + """) + +def _create_link_agencies_locations_table(): + op.execute(""" + create table if not exists public.link_agencies_locations + ( + id serial + primary key, + agency_id integer not null + references public.agencies + on delete cascade, + location_id integer not null + references public.locations + on delete cascade, + constraint unique_agency_location + unique (agency_id, location_id) + ) + """) + +def _remove_link_agencies_locations_table(): + op.drop_table(LINK_AGENCIES_LOCATIONS_TABLE_NAME) + +def _remove_locations_table(): + op.drop_table(LOCATIONS_TABLE_NAME) + +def _remove_localities_table(): + op.drop_table(LOCALITIES_TABLE_NAME) + +def _remove_counties_table(): + op.drop_table(COUNTIES_TABLE_NAME) + +def _remove_us_states_table(): + op.drop_table(US_STATES_TABLE_NAME) diff --git a/src/db/models/helpers.py b/src/db/models/helpers.py index e4b941ed..1782b1e9 100644 --- a/src/db/models/helpers.py +++ b/src/db/models/helpers.py @@ -40,4 +40,25 @@ def url_id_column() -> Column[int]: CURRENT_TIME_SERVER_DEFAULT = func.now() def url_id_primary_key_constraint() -> PrimaryKeyConstraint: - return PrimaryKeyConstraint('url_id') \ No newline at end of file + return PrimaryKeyConstraint('url_id') + +def county_column(nullable: bool = False) -> Column[int]: + return Column( + Integer(), + ForeignKey('counties.id', ondelete='CASCADE'), + nullable=nullable + ) + +def locality_column(nullable: bool = False) -> Column[int]: + return Column( + Integer(), + ForeignKey('localities.id', ondelete='CASCADE'), + nullable=nullable + ) + +def us_state_column(nullable: bool = False) -> Column[int]: + return Column( + Integer(), + ForeignKey('us_states.id', ondelete='CASCADE'), + nullable=nullable + ) \ No newline at end of file diff --git a/src/db/models/impl/link/agency_location/__init__.py b/src/db/models/impl/link/agency_location/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/db/models/impl/link/agency_location/sqlalchemy.py b/src/db/models/impl/link/agency_location/sqlalchemy.py new file mode 100644 index 00000000..18a3ae5f --- /dev/null +++ b/src/db/models/impl/link/agency_location/sqlalchemy.py @@ -0,0 +1,10 @@ +from src.db.models.mixins import AgencyDependentMixin, LocationDependentMixin +from src.db.models.templates_.with_id import WithIDBase + + +class LinkAgencyLocation( + WithIDBase, + AgencyDependentMixin, + LocationDependentMixin, +): + __tablename__ = "link_agencies_locations" \ No newline at end of file diff --git a/src/db/models/impl/location/__init__.py b/src/db/models/impl/location/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/db/models/impl/location/county/__init__.py b/src/db/models/impl/location/county/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/db/models/impl/location/county/sqlalchemy.py b/src/db/models/impl/location/county/sqlalchemy.py new file mode 100644 index 00000000..b3428449 --- /dev/null +++ b/src/db/models/impl/location/county/sqlalchemy.py @@ -0,0 +1,18 @@ +from sqlalchemy import String, Column, Float, Integer +from sqlalchemy.orm import Mapped + +from src.db.models.helpers import us_state_column +from src.db.models.templates_.with_id import WithIDBase + + +class County( + WithIDBase, +): + __tablename__ = "counties" + + name: Mapped[str] + state_id = us_state_column() + fips: Mapped[str] = Column(String(5), nullable=True) + lat: Mapped[float] = Column(Float, nullable=True) + lng: Mapped[float] = Column(Float, nullable=True) + population: Mapped[int] = Column(Integer, nullable=True) \ No newline at end of file diff --git a/src/db/models/impl/location/locality/__init__.py b/src/db/models/impl/location/locality/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/db/models/impl/location/locality/sqlalchemy.py b/src/db/models/impl/location/locality/sqlalchemy.py new file mode 100644 index 00000000..216706fd --- /dev/null +++ b/src/db/models/impl/location/locality/sqlalchemy.py @@ -0,0 +1,14 @@ +from sqlalchemy import String, Column + +from src.db.models.helpers import county_column +from src.db.models.templates_.with_id import WithIDBase + + +class Locality( + WithIDBase, +): + + __tablename__ = "localities" + + name = Column(String(255), nullable=False) + county_id = county_column(nullable=False) diff --git a/src/db/models/impl/location/location/__init__.py b/src/db/models/impl/location/location/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/db/models/impl/location/location/enums.py b/src/db/models/impl/location/location/enums.py new file mode 100644 index 00000000..24a99ce9 --- /dev/null +++ b/src/db/models/impl/location/location/enums.py @@ -0,0 +1,8 @@ +from enum import Enum + + +class LocationType(Enum): + NATIONAL = "National" + STATE = "State" + COUNTY = "County" + LOCALITY = "Locality" \ No newline at end of file diff --git a/src/db/models/impl/location/location/sqlalchemy.py b/src/db/models/impl/location/location/sqlalchemy.py new file mode 100644 index 00000000..1a5dc435 --- /dev/null +++ b/src/db/models/impl/location/location/sqlalchemy.py @@ -0,0 +1,19 @@ +from sqlalchemy import Float, Column + +from src.db.models.helpers import us_state_column, county_column, locality_column, enum_column +from src.db.models.impl.location.location.enums import LocationType +from src.db.models.templates_.with_id import WithIDBase + + +class Location( + WithIDBase +): + + __tablename__ = "locations" + + state_id = us_state_column(nullable=True) + county_id = county_column(nullable=True) + locality_id = locality_column(nullable=True) + type = enum_column(LocationType, name="location_type", nullable=False) + lat = Column(Float(), nullable=True) + lng = Column(Float(), nullable=True) diff --git a/src/db/models/impl/location/us_state/__init__.py b/src/db/models/impl/location/us_state/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/db/models/impl/location/us_state/sqlalchemy.py b/src/db/models/impl/location/us_state/sqlalchemy.py new file mode 100644 index 00000000..c4cdfc2f --- /dev/null +++ b/src/db/models/impl/location/us_state/sqlalchemy.py @@ -0,0 +1,12 @@ +from sqlalchemy.orm import Mapped + +from src.db.models.templates_.with_id import WithIDBase + + +class USState( + WithIDBase, +): + __tablename__ = "us_states" + + state_name: Mapped[str] + state_iso: Mapped[str] \ No newline at end of file diff --git a/src/db/models/mixins.py b/src/db/models/mixins.py index d0dbbcab..12a0b2a1 100644 --- a/src/db/models/mixins.py +++ b/src/db/models/mixins.py @@ -38,6 +38,15 @@ class BatchDependentMixin: nullable=False ) +class LocationDependentMixin: + location_id = Column( + Integer, + ForeignKey( + 'locations.id', + ondelete="CASCADE", + ), + nullable=False + ) class AgencyDependentMixin: agency_id = Column(