add region to database (#42)

* add region to database

* make function names and help more descriptive

backend/src/mirrors_qa_backend/cli/country.py

@@ -0,0 +1,53 @@
+import csv
+
+from mirrors_qa_backend import logger
+from mirrors_qa_backend.db import Session
+from mirrors_qa_backend.db.country import create_country
+from mirrors_qa_backend.db.region import create_region
+from mirrors_qa_backend.schemas import Country, Region
+
+
+def create_regions_and_countries(countries: list[Country]) -> None:
+    """Create the region and associated countries in the database."""
+    with Session.begin() as session:
+        for country in countries:
+            db_country = create_country(
+                session,
+                country_code=country.code,
+                country_name=country.name,
+            )
+            if country.region:
+                db_region = create_region(
+                    session,
+                    region_code=country.region.code,
+                    region_name=country.region.name,
+                )
+                db_country.region = db_region
+                session.add(db_country)
+
+
+def extract_country_regions_from_csv(csv_data: list[str]) -> list[Country]:
+    regions: list[Country] = []
+    for row in csv.DictReader(csv_data):
+        country_code = row["country_iso_code"]
+        country_name = row["country_name"]
+        region_code = row["continent_code"]
+        region_name = row["continent_name"]
+        if all([country_code, country_name, region_code, region_name]):
+            regions.append(
+                Country(
+                    code=country_code.lower(),
+                    name=country_name.title(),
+                    region=Region(
+                        code=region_code.lower(),
+                        name=region_name.title(),
+                    ),
+                )
+            )
+        else:
+            logger.critical(
+                f"Skipping row with missing entries: country_code: {country_code}, "
+                f"country_name: {country_name}, region_code: {region_code}, "
+                f"region_name: {region_name}"
+            )
+    return regions

backend/src/mirrors_qa_backend/db/mirrors.py

@@ -4,8 +4,10 @@
 from sqlalchemy.orm import Session as OrmSession
 
 from mirrors_qa_backend import logger, schemas
+from mirrors_qa_backend.db.country import get_country_or_none
 from mirrors_qa_backend.db.exceptions import EmptyMirrorsError, RecordDoesNotExistError
 from mirrors_qa_backend.db.models import Mirror
+from mirrors_qa_backend.db.region import get_region_or_none
 
 
 @dataclass
@@ -16,6 +18,17 @@ class MirrorsUpdateResult:
     nb_mirrors_disabled: int = 0
 
 
+def update_mirror_country(
+    session: OrmSession, country_code: str, mirror: Mirror
+) -> Mirror:
+    logger.debug("Updating mirror country information.")
+    mirror.country = get_country_or_none(session, country_code)
+    if mirror.country and mirror.country.region_code:
+        mirror.region = get_region_or_none(session, mirror.country.region_code)
+    session.add(mirror)
+    return mirror
+
+
 def create_mirrors(session: OrmSession, mirrors: list[schemas.Mirror]) -> int:
     """Number of mirrors created in the database.
 
@@ -27,7 +40,6 @@ def create_mirrors(session: OrmSession, mirrors: list[schemas.Mirror]) -> int:
             id=mirror.id,
             base_url=mirror.base_url,
             enabled=mirror.enabled,
-            region=mirror.region,
             asn=mirror.asn,
             score=mirror.score,
             latitude=mirror.latitude,
@@ -37,7 +49,12 @@ def create_mirrors(session: OrmSession, mirrors: list[schemas.Mirror]) -> int:
             as_only=mirror.as_only,
             other_countries=mirror.other_countries,
         )
+
         session.add(db_mirror)
+
+        if mirror.country_code:
+            update_mirror_country(session, mirror.country_code, db_mirror)
+
         logger.debug(f"Registered new mirror: {db_mirror.id}.")
         nb_created += 1
     return nb_created
@@ -90,6 +107,13 @@ def create_or_update_mirror_status(
             db_mirror.enabled = True
             session.add(db_mirror)
             result.nb_mirrors_added += 1
+
+        # New mirrors DB model contain country data. As such, we update the
+        # country information regardless of the status update.
+        if db_mirror_id in current_mirrors:
+            country_code = current_mirrors[db_mirror_id].country_code
+            if country_code:
+                update_mirror_country(session, country_code, db_mirror)
     return result
 
 

backend/src/mirrors_qa_backend/db/models.py

@@ -64,6 +64,21 @@ class WorkerCountry(Base):
     )
 
 
+class Region(Base):
+    """Continental region."""
+
+    __tablename__ = "region"
+
+    code: Mapped[str] = mapped_column(primary_key=True)  # continent code
+    name: Mapped[str]  # continent name
+    countries: Mapped[list[Country]] = relationship(
+        back_populates="region", init=False, repr=False
+    )
+    mirrors: Mapped[list[Mirror]] = relationship(
+        back_populates="region", init=False, repr=False
+    )
+
+
 class Country(Base):
     """Country where a worker runs tests for a mirror."""
 
@@ -74,6 +89,17 @@ class Country(Base):
     )  # two-letter country codes as defined in ISO 3166-1
 
     name: Mapped[str]  # full name of the country (in English)
+    region_code: Mapped[str | None] = mapped_column(
+        ForeignKey("region.code"), init=False, default=None
+    )
+
+    region: Mapped[Region | None] = relationship(
+        back_populates="countries", init=False, repr=False
+    )
+
+    mirrors: Mapped[list[Mirror]] = relationship(
+        back_populates="country", init=False, repr=False
+    )
 
     workers: Mapped[list[Worker]] = relationship(
         back_populates="countries",
@@ -91,8 +117,13 @@ class Mirror(Base):
     id: Mapped[str] = mapped_column(primary_key=True)  # hostname of a mirror URL
     base_url: Mapped[str]
     enabled: Mapped[bool]
+    region_code: Mapped[str | None] = mapped_column(
+        ForeignKey("region.code"), init=False, default=None
+    )
+    country_code: Mapped[str | None] = mapped_column(
+        ForeignKey("country.code"), init=False, default=None
+    )
     # metadata of a mirror from MirroBrain (https://mirrorbrain-docs.readthedocs.io/en/latest/mirrors.html#displaying-details-about-a-mirror)
-    region: Mapped[str | None] = mapped_column(default=None)
     asn: Mapped[str | None] = mapped_column(default=None)
     score: Mapped[int | None] = mapped_column(default=None)
     latitude: Mapped[float | None] = mapped_column(default=None)
@@ -106,6 +137,14 @@ class Mirror(Base):
         back_populates="mirror", init=False, repr=False
     )
 
+    country: Mapped[Country | None] = relationship(
+        back_populates="mirrors", init=False, repr=False
+    )
+
+    region: Mapped[Region | None] = relationship(
+        back_populates="mirrors", init=False, repr=False
+    )
+
     __table_args__ = (UniqueConstraint("base_url"),)
 
 

backend/src/mirrors_qa_backend/db/region.py

@@ -0,0 +1,36 @@
+from sqlalchemy import select
+from sqlalchemy.dialects.postgresql import insert
+from sqlalchemy.orm import Session as OrmSession
+
+from mirrors_qa_backend.db.exceptions import RecordDoesNotExistError
+from mirrors_qa_backend.db.models import Country, Region
+
+
+def get_countries_for(session: OrmSession, region_code: str) -> list[Country]:
+    """Get countries belonging to the provided region."""
+
+    return list(
+        session.scalars(select(Country).where(Country.region_code == region_code)).all()
+    )
+
+
+def get_region_or_none(session: OrmSession, region_code: str) -> Region | None:
+    return session.scalars(
+        select(Region).where(Region.code == region_code)
+    ).one_or_none()
+
+
+def get_region(session: OrmSession, region_code: str) -> Region:
+    if region := get_region_or_none(session, region_code):
+        return region
+    raise RecordDoesNotExistError(f"Region with code {region_code} does not exist.")
+
+
+def create_region(session: OrmSession, *, region_code: str, region_name: str) -> Region:
+    """Creates a new continental region in the database."""
+    session.execute(
+        insert(Region)
+        .values(code=region_code, name=region_name)
+        .on_conflict_do_nothing(index_elements=["code"])
+    )
+    return get_region(session, region_code)

backend/src/mirrors_qa_backend/entrypoint.py

@@ -6,6 +6,10 @@
 
 from mirrors_qa_backend import logger
 from mirrors_qa_backend.__about__ import __version__
+from mirrors_qa_backend.cli.country import (
+    create_regions_and_countries,
+    extract_country_regions_from_csv,
+)
 from mirrors_qa_backend.cli.mirrors import update_mirrors
 from mirrors_qa_backend.cli.scheduler import main as start_scheduler
 from mirrors_qa_backend.cli.worker import create_worker, update_worker
@@ -15,6 +19,7 @@
 CREATE_WORKER_CLI = "create-worker"
 UPDATE_WORKER_CLI = "update-worker"
 SCHEDULER_CLI = "scheduler"
+CREATE_COUNTRY_REGIONS_CLI = "create-countries"
 
 
 def main():
@@ -95,6 +100,21 @@ def main():
         UPDATE_WORKER_CLI, help="Update a worker", parents=[worker_parser]
     )
 
+    create_country_regions_cli = subparsers.add_parser(
+        CREATE_COUNTRY_REGIONS_CLI, help="Create countries and associated regions."
+    )
+    create_country_regions_cli.add_argument(
+        "country_region_csv_file",
+        metavar="csv-file",
+        type=argparse.FileType("r", encoding="utf-8"),
+        nargs="?",
+        default=sys.stdin,
+        help=(
+            "CSV file containing countries and associated regions "
+            "(format: Maxmind's GeoIPLite Country Locations csv) (default: stdin)."
+        ),
+    )
+
     args = parser.parse_args()
     if args.verbose:
         logger.setLevel(logging.DEBUG)
@@ -137,6 +157,19 @@ def main():
             logger.error(f"error while updating worker: {exc!s}")
             sys.exit(1)
         logger.info(f"Updated countries for worker {args.worker_id!r}")
+    elif args.cli_name == CREATE_COUNTRY_REGIONS_CLI:
+        try:
+            logger.debug("Creating regions and associated countries.")
+
+            create_regions_and_countries(
+                extract_country_regions_from_csv(
+                    args.country_region_csv_file.readlines()
+                )
+            )
+        except Exception as exc:
+            logger.error(f"error while creating regions: {exc!s}")
+            sys.exit(1)
+        logger.info("Created regions and associated countries.")
     else:
         args.print_help()
 

backend/src/mirrors_qa_backend/extract.py

@@ -51,13 +51,15 @@ def is_country_row(tag: Tag) -> bool:
         hostname: Any = urlsplit(
             base_url
         ).netloc  # pyright: ignore [reportUnknownMemberType]
+        country_code = row.find("img")["alt"].lower()
         if hostname in Settings.MIRRORS_EXCLUSION_LIST:
             continue
         mirrors.append(
             schemas.Mirror(
                 id=hostname,
                 base_url=base_url,
                 enabled=True,
+                country_code=country_code,
             )
         )
     return mirrors

backend/src/mirrors_qa_backend/migrations/versions/074ae280bb70_introduce_regions.py

@@ -0,0 +1,73 @@
+"""introduce regions
+
+Revision ID: 074ae280bb70
+Revises: 17d587447299
+Create Date: 2024-08-22 11:57:17.239215
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "074ae280bb70"
+down_revision = "17d587447299"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "region",
+        sa.Column("code", sa.String(), nullable=False),
+        sa.Column("name", sa.String(), nullable=False),
+        sa.PrimaryKeyConstraint("code", name=op.f("pk_region")),
+    )
+    op.add_column("country", sa.Column("region_code", sa.String(), nullable=True))
+    op.create_foreign_key(
+        op.f("fk_country_region_code_region"),
+        "country",
+        "region",
+        ["region_code"],
+        ["code"],
+    )
+    op.add_column("mirror", sa.Column("region_code", sa.String(), nullable=True))
+    op.add_column("mirror", sa.Column("country_code", sa.String(), nullable=True))
+    op.create_foreign_key(
+        op.f("fk_mirror_country_code_country"),
+        "mirror",
+        "country",
+        ["country_code"],
+        ["code"],
+    )
+    op.create_foreign_key(
+        op.f("fk_mirror_region_code_region"),
+        "mirror",
+        "region",
+        ["region_code"],
+        ["code"],
+    )
+    op.drop_column("mirror", "region")
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column(
+        "mirror", sa.Column("region", sa.VARCHAR(), autoincrement=False, nullable=True)
+    )
+    op.drop_constraint(
+        op.f("fk_mirror_region_code_region"), "mirror", type_="foreignkey"
+    )
+    op.drop_constraint(
+        op.f("fk_mirror_country_code_country"), "mirror", type_="foreignkey"
+    )
+    op.drop_column("mirror", "country_code")
+    op.drop_column("mirror", "region_code")
+    op.drop_constraint(
+        op.f("fk_country_region_code_region"), "country", type_="foreignkey"
+    )
+    op.drop_column("country", "region_code")
+    op.drop_table("region")
+    # ### end Alembic commands ###