loctrkd/ocid_dload.py

   1 from configparser import ConfigParser, NoOptionError
   2 import csv
   3 from logging import getLogger
   4 import requests
   5 from sqlite3 import connect
   6 from typing import Any, IO, Optional
   7 from zlib import decompressobj, MAX_WBITS
   8
   9 from . import common
  10
  11 log = getLogger("loctrkd/ocid_dload")
  12
  13 RURL = (
  14     "https://opencellid.org/ocid/downloads"
  15     "?token={token}&type={dltype}&file={fname}.csv.gz"
  16 )
  17
  18 SCHEMA = """create table if not exists cells (
  19   "radio" text,
  20   "mcc" int,
  21   "net" int,
  22   "area" int,
  23   "cell" int,
  24   "unit" int,
  25   "lon" int,
  26   "lat" int,
  27   "range" int,
  28   "samples" int,
  29   "changeable" int,
  30   "created" int,
  31   "updated" int,
  32   "averageSignal" int
  33 )"""
  34 DBINDEX = "create index if not exists cell_idx on cells (area, cell)"
  35
  36
  37 class unzipped:
  38     """
  39     File-like object that unzips http response body.
  40     read(size) method returns chunks of binary data as bytes
  41     When used as iterator, splits data to lines
  42     and yelds them as strings.
  43     """
  44
  45     def __init__(self, zstream: IO[bytes]) -> None:
  46         self.zstream = zstream
  47         self.decoder: Optional[Any] = decompressobj(16 + MAX_WBITS)
  48         self.outdata = b""
  49         self.line = b""
  50
  51     def read(self, n: int = 1024) -> bytes:
  52         if self.decoder is None:
  53             return b""
  54         while len(self.outdata) < n:
  55             raw_data = self.zstream.read(n)
  56             self.outdata += self.decoder.decompress(raw_data)
  57             if not raw_data:
  58                 self.decoder = None
  59                 break
  60         if self.outdata:
  61             data, self.outdata = self.outdata[:n], self.outdata[n:]
  62             return data
  63         return b""
  64
  65     def __next__(self) -> str:
  66         while True:
  67             splittry = self.line.split(b"\n", maxsplit=1)
  68             if len(splittry) > 1:
  69                 break
  70             moredata = self.read(256)
  71             if not moredata:
  72                 raise StopIteration
  73             self.line += moredata
  74         line, rest = splittry
  75         self.line = rest
  76         return line.decode("utf-8")
  77
  78     def __iter__(self) -> "unzipped":
  79         return self
  80
  81
  82 def main(conf: ConfigParser) -> None:
  83     try:
  84         url = conf.get("opencellid", "downloadurl")
  85         mcc = "<unspecified>"
  86     except NoOptionError:
  87         try:
  88             with open(
  89                 conf.get("opencellid", "downloadtoken"), encoding="ascii"
  90             ) as fl:
  91                 token = fl.read().strip()
  92         except FileNotFoundError:
  93             log.warning(
  94                 "Opencellid access token not configured, cannot download"
  95             )
  96             return
  97         mcc = conf.get("opencellid", "downloadmcc")
  98         if mcc == "full":
  99             dltype = "full"
 100             fname = "cell_towers"
 101         else:
 102             dltype = "mcc"
 103             fname = mcc
 104         url = RURL.format(token=token, dltype="mcc", fname=mcc)
 105     dbfn = conf.get("opencellid", "dbfn")
 106     count = 0
 107     with requests.get(url, stream=True) as resp, connect(dbfn) as db:
 108         log.debug("Requested %s, result %s", url, resp)
 109         if resp.status_code != 200:
 110             log.error("Error getting %s: %s", url, resp)
 111             return
 112         db.execute("pragma journal_mode = wal")
 113         db.execute(SCHEMA)
 114         db.execute("delete from cells")
 115         rows = csv.reader(unzipped(resp.raw))
 116         for row in rows:
 117             db.execute(
 118                 """insert into cells
 119                    values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
 120                 row,
 121             )
 122             count += 1
 123         if count < 1:
 124             db.rollback()
 125             log.warning("Did not get any data for MCC %s, rollback", mcc)
 126         else:
 127             db.execute(DBINDEX)
 128             db.commit()
 129             log.info(
 130                 "repopulated %s with %d records for MCC %s", dbfn, count, mcc
 131             )
 132
 133
 134 if __name__.endswith("__main__"):
 135     main(common.init(log))