Bg geoip update (#951)

* Add checkin models and initial route

* Porting old code to new probe_services endpoint

* Add geoip2 dependency

* Add geoip tables as dependencies

* Ported probe_geoip function

* Porting more functions from checkin

* Added health checks for clickhouse

* First complete version of checkin endpoint port

* Fix bad name for metric

* Port some checkin integration tests

* Add pytest-docker as test dependency

* Add docker compose to launch clickhouse db for tests

* Making tests work with the new setup

* Fixed unhashable settings error

* ported script to download geoip DB

* Remove exit

* Use setting to fetch db dir

* Added setting for geoip db dir location

* Rename download geoip file

* Set download dir as input argument

* download db before tests

* Black reformat

* fix error with bad header parsing

* Making tests work

* Add settings as a parameter to mock it in tests

* Fix bad reader dependency

* Skip broken test

* Add function to populate db

* Add data for tests

* Fix fixture loading

* Add clickhouse db fixture

* Adding integration tests for check in

* Removed todo comments

* Add missin / at the start of path

* Black reformat

* Refactor to move prio code into a prio.py file

* Add todo comment

* Add timer to measure function run time

* Improved writing of comment

* Improved writing of comment

* Added missing alembic migration

* Revert "Added missing alembic migration"

This reverts commit fb072d2.

* Add help comment in init.sh

* Black reformat

* Add APScheduler for background tasks

* Add geoip db update job

* Changed APScheduler for fast api utils repeat-every

* Add repeating task for updating the geoip DB

* Add parameter to enable/disable repeating tasks

The tests were hanging because of the repeating tasks, so I added
a parameter to deactivate them in tests

* Change task execution frequency

* Remove print

* Black reformat

* Add TS calculation to function body

* Re-raise consistency check exception in download_geoip

* Add geoip data dir fixture

* Removed unnecessary remove

* Changed heavy fastpath query on health check

* trigger ci

* Reformat query

* Changed health check to a less heavy query

Author: LDiazN

ooniapi/common/src/common/config.py

@@ -31,3 +31,6 @@ class Settings(BaseSettings):
     email_source_address: str = "[email protected]"
 
     vpn_credential_refresh_hours: int = 24
+
+    # Where the geoip DBs are downloaded to
+    geoip_db_dir: str = "/var/lib/ooni/geoip"

ooniapi/common/src/common/metrics.py

@@ -1,6 +1,7 @@
 import secrets
 
-from typing import Annotated
+from typing import Annotated, Optional
+import timeit
 
 from fastapi import FastAPI, Response, Depends, HTTPException, status
 from fastapi.security import HTTPBasic, HTTPBasicCredentials
@@ -10,6 +11,7 @@
     CONTENT_TYPE_LATEST,
     CollectorRegistry,
     generate_latest,
+    Histogram,
 )
 
 security = HTTPBasic()
@@ -40,3 +42,51 @@ def metrics(
 
     endpoint = "/metrics"
     app.get(endpoint, include_in_schema=True, tags=None)(metrics)
+
+
+FUNCTION_TIME = Histogram(
+    "fn_execution_time_seconds",
+    "Function execution time in seconds",
+    labelnames=["name", "status"],
+)
+
+
+def timer(func, timer_name: Optional[str] = None):
+    """Measure function execution time in seconds.
+
+    Metric will include execution status (error or success) and function name, you
+    can also specify another name with `timer_name`. Example:
+
+    ```
+    @timer
+    def add(x,y):
+        return x + y
+    ```
+
+    Args:
+        func (Callable): function to time
+        timer_name (Optional[str], optional): Alternative name for this timer, uses function name if
+        not provided. Defaults to None.
+    """
+
+    name = timer_name or func.__name__
+
+    def timed_function(*args, **kwargs):
+
+        start_time = timeit.default_timer()
+        result = None
+        successs = False
+
+        try:
+            result = func(*args, **kwargs)
+            successs = True
+        finally:
+            end_time = timeit.default_timer()
+            status = "success" if successs else "error"
+            FUNCTION_TIME.labels(name=name, status=status).observe(
+                end_time - start_time
+            )
+
+        return result
+
+    return timed_function

ooniapi/services/oonimeasurements/src/oonimeasurements/main.py

@@ -85,8 +85,10 @@ async def health(
     errors = []
 
     try:
-        query = """SELECT *
-        FROM fastpath FINAL
+        query = """
+        SELECT COUNT()
+        FROM fastpath
+        WHERE measurement_start_time < NOW() AND measurement_start_time > NOW() - INTERVAL 3 HOUR
         """
         query_click(db=db, query=query, query_params={})
     except Exception as exc:

ooniapi/services/ooniprobe/pyproject.toml

@@ -24,6 +24,8 @@ dependencies = [
   "prometheus-fastapi-instrumentator ~= 6.1.0",
   "prometheus-client",
   "pem ~= 23.1.0",
+  "geoip2 ~= 5.0.1",
+  "fastapi-utils[all] ~= 0.8.0"
 ]
 
 readme = "README.md"
@@ -67,6 +69,7 @@ dependencies = [
   "pytest-postgresql",
   "pytest-asyncio",
   "freezegun",
+  "pytest-docker",
 ]
 path = ".venv/"
 

ooniapi/services/ooniprobe/src/ooniprobe/dependencies.py

@@ -1,16 +1,23 @@
-from functools import lru_cache
-from typing import Annotated
+from typing import Annotated, TypeAlias
+from pathlib import Path
 
 from fastapi import Depends
 
+import geoip2.database
+
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
 
+from clickhouse_driver import Client as Clickhouse
+
 from .common.config import Settings
 from .common.dependencies import get_settings
 
 
-def get_postgresql_session(settings: Annotated[Settings, Depends(get_settings)]):
+SettingsDep: TypeAlias = Annotated[Settings, Depends(get_settings)]
+
+
+def get_postgresql_session(settings: SettingsDep):
     engine = create_engine(settings.postgresql_url)
     SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
 
@@ -19,3 +26,30 @@ def get_postgresql_session(settings: Annotated[Settings, Depends(get_settings)])
         yield db
     finally:
         db.close()
+
+
+def get_cc_reader(settings: SettingsDep):
+    db_path = Path(settings.geoip_db_dir, "cc.mmdb")
+    return geoip2.database.Reader(db_path)
+
+
+CCReaderDep = Annotated[geoip2.database.Reader, Depends(get_cc_reader)]
+
+
+def get_asn_reader(settings: SettingsDep):
+    db_path = Path(settings.geoip_db_dir, "asn.mmdb")
+    return geoip2.database.Reader(db_path)
+
+
+ASNReaderDep = Annotated[geoip2.database.Reader, Depends(get_asn_reader)]
+
+
+def get_clickhouse_session(settings: SettingsDep):
+    db = Clickhouse.from_url(settings.clickhouse_url)
+    try:
+        yield db
+    finally:
+        db.disconnect()
+
+
+ClickhouseDep = Annotated[Clickhouse, Depends(get_clickhouse_session)]

ooniapi/services/ooniprobe/src/ooniprobe/download_geoip.py

@@ -0,0 +1,145 @@
+"""
+Updates asn.mmdb and cc.mmdb in /var/lib/ooniapi/
+
+"""
+
+import sys
+import gzip
+import timeit
+import shutil
+import logging
+
+import geoip2.database
+from pathlib import Path
+from datetime import datetime, timezone
+from urllib.error import HTTPError
+from urllib.request import urlopen, Request
+
+from prometheus_client import metrics
+
+
+class Metrics:
+    GEOIP_ASN_NODE_CNT = metrics.Gauge("geoip_asn_node_cnt", "Count of geoi nodes")
+    GEOIP_ASN_EPOCH = metrics.Gauge("geoip_asn_epoch", "Geoip current ASN epoch")
+    GEOIP_CC_NODE_CNT = metrics.Gauge("geoip_cc_node_cnt", "Geoip asn node count")
+    GEOIP_CC_EPOCH = metrics.Gauge("geoip_cc_epoch", "Geoip current CC epoch")
+    GEOIP_CHECKFAIL = metrics.Counter(
+        "ooni_geoip_checkfail", "How many times did the check fail in geo ip fail"
+    )
+    GEOIP_UPDATED = metrics.Counter(
+        "ooni_geoip_updated", "How many times was the geoip database updated"
+    )
+    GEOIP_DOWNLOAD_TIME = metrics.Histogram(
+        "geoip_download_time", "How long it takes to download the DB"
+    )
+
+log = logging.getLogger("ooni_download_geoip")
+
+log.addHandler(logging.StreamHandler(sys.stdout))
+log.setLevel(logging.DEBUG)
+
+
+def get_request(url):
+    req = Request(url)
+    # We need to set the user-agent otherwise db-ip gives us a 403
+    req.add_header("User-Agent", "ooni-downloader")
+    return urlopen(req)
+
+
+def is_already_updated(db_dir: Path, ts : str) -> bool:
+    try:
+        with (db_dir / "geoipdbts").open() as in_file:
+            current_ts = in_file.read()
+    except FileNotFoundError:
+        return False
+
+    return current_ts == ts
+
+
+def is_latest_available(url: str) -> bool:
+    log.info(f"fetching {url}")
+    try:
+        resp = get_request(url)
+        return resp.status == 200
+    except HTTPError as err:
+        if resp.status == 404:  # type: ignore
+            log.info(f"{url} hasn't been updated yet")
+            return False
+        log.info(f"unexpected status code '{err.code}' in {url}")
+        return False
+
+
+def check_geoip_db(path: Path) -> None:
+    assert "cc" in path.name or "asn" in path.name, "invalid path"
+
+    with geoip2.database.Reader(str(path)) as reader:
+        if "asn" in path.name:
+            r1 = reader.asn("8.8.8.8")
+            assert r1 is not None, "database file is invalid"
+            m = reader.metadata()
+            Metrics.GEOIP_ASN_NODE_CNT.set(m.node_count)
+            Metrics.GEOIP_ASN_EPOCH.set(m.build_epoch)
+
+        elif "cc" in path.name:
+            r2 = reader.country("8.8.8.8")
+            assert r2 is not None, "database file is invalid"
+            m = reader.metadata()
+            Metrics.GEOIP_CC_NODE_CNT.set(m.node_count)
+            Metrics.GEOIP_CC_EPOCH.set(m.build_epoch)
+
+
+def download_geoip(db_dir: Path, url: str, filename: str) -> None:
+    start_time = timeit.default_timer()  # Start timer
+    log.info(f"Updating geoip database for {url} ({filename})")
+
+    tmp_gz_out = db_dir / f"{filename}.gz.tmp"
+    tmp_out = db_dir / f"{filename}.tmp"
+
+    with get_request(url) as resp:
+        with tmp_gz_out.open("wb") as out_file:
+            shutil.copyfileobj(resp, out_file)
+    with gzip.open(str(tmp_gz_out)) as in_file:
+        with tmp_out.open("wb") as out_file:
+            shutil.copyfileobj(in_file, out_file)
+    tmp_gz_out.unlink()
+
+    try:
+        check_geoip_db(tmp_out)
+    except Exception as exc:
+        log.error(f"consistenty check on the geoip DB failed: {exc}")
+        Metrics.GEOIP_CHECKFAIL.inc()
+        raise 
+
+    tmp_out.rename(db_dir / filename)
+    endtime = timeit.default_timer()  # End timer
+    Metrics.GEOIP_DOWNLOAD_TIME.observe(endtime - start_time)
+
+
+def update_geoip(db_dir: Path, ts : str, asn_url : str, cc_url : str) -> None:
+    db_dir.mkdir(parents=True, exist_ok=True)
+    download_geoip(db_dir, asn_url, "asn.mmdb")
+    download_geoip(db_dir, cc_url, "cc.mmdb")
+
+    with (db_dir / "geoipdbts").open("w") as out_file:
+        out_file.write(ts)
+
+    log.info("Updated GeoIP databases")
+    Metrics.GEOIP_UPDATED.inc()
+
+
+def try_update(db_dir: str):
+    db_dir_path = Path(db_dir)
+
+    ts = datetime.now(timezone.utc).strftime("%Y-%m")
+    asn_url = f"https://download.db-ip.com/free/dbip-asn-lite-{ts}.mmdb.gz"
+    cc_url = f"https://download.db-ip.com/free/dbip-country-lite-{ts}.mmdb.gz"
+
+    if is_already_updated(db_dir_path, ts):
+        log.debug("Database already updated. Exiting.")
+        return
+
+    if not is_latest_available(asn_url) or not is_latest_available(cc_url):
+        log.debug("Update not available yet. Exiting.")
+        return
+
+    update_geoip(db_dir_path, ts, asn_url, cc_url)