229 lines
7.2 KiB
Python
Executable file
229 lines
7.2 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
from pathlib import Path
|
|
import functools
|
|
from datetime import datetime
|
|
from itertools import islice
|
|
from typing import Type, NamedTuple, Union, Optional
|
|
import logging
|
|
|
|
from location import _load_locations, Location, get_logger
|
|
import sqlalchemy # type: ignore
|
|
import sqlalchemy as sa # type: ignore
|
|
|
|
from kython import ichunks
|
|
|
|
|
|
from kython.py37 import fromisoformat
|
|
|
|
# TODO move to some common thing?
|
|
class IsoDateTime(sqlalchemy.TypeDecorator):
|
|
# in theory could use something more effecient? e.g. blob for encoded datetime and tz?
|
|
# but practically, the difference seems to be pretty small, so perhaps fine for now
|
|
impl = sqlalchemy.types.String
|
|
|
|
# TODO optional?
|
|
def process_bind_param(self, value: Optional[datetime], dialect) -> Optional[str]:
|
|
if value is None:
|
|
return None
|
|
return value.isoformat()
|
|
|
|
def process_result_value(self, value: Optional[str], dialect) -> Optional[datetime]:
|
|
if value is None:
|
|
return None
|
|
return fromisoformat(value)
|
|
|
|
|
|
def _map_type(cls):
|
|
tmap = {
|
|
str: sa.String,
|
|
float: sa.Float,
|
|
datetime: IsoDateTime,
|
|
}
|
|
r = tmap.get(cls, None)
|
|
if r is not None:
|
|
return r
|
|
|
|
|
|
if getattr(cls, '__origin__', None) == Union:
|
|
elems = cls.__args__
|
|
elems = [e for e in elems if e != type(None)]
|
|
if len(elems) == 1:
|
|
return _map_type(elems[0]) # meh..
|
|
raise RuntimeError(f'Unexpected type {cls}')
|
|
|
|
# TODO to strart with, just assert utc when serializing, deserializing
|
|
# TODO how to use timestamp as key? just round it?
|
|
|
|
def make_schema(cls: Type[NamedTuple]): # TODO covariant?
|
|
res = []
|
|
for name, ann in cls.__annotations__.items():
|
|
res.append(sa.Column(name, _map_type(ann)))
|
|
return res
|
|
|
|
|
|
def get_table(db_path: Path, type_, name='table'):
|
|
db = sa.create_engine(f'sqlite:///{db_path}')
|
|
engine = db.connect() # TODO do I need to tear anything down??
|
|
meta = sa.MetaData(engine)
|
|
schema = make_schema(type_)
|
|
sa.Table(name, meta, *schema)
|
|
meta.create_all()
|
|
table = sa.table(name, *schema)
|
|
return engine, table
|
|
|
|
def cache_locs(source: Path, db_path: Path, limit=None):
|
|
engine, table = get_table(db_path=db_path, type_=Location)
|
|
|
|
with source.open('r') as fo:
|
|
# TODO fuck. do I really need to split myself??
|
|
# sqlalchemy.exc.OperationalError: (sqlite3.OperationalError) too many SQL variables
|
|
# TODO count deprecated??
|
|
# print(engine.execute(table.count()).fetchone())
|
|
for chunk in ichunks(islice(_load_locations(fo), 0, limit), 10000):
|
|
engine.execute(table.insert().values(chunk))
|
|
|
|
# TODO maintain order during insertion?
|
|
|
|
def iter_db_locs(db_path: Path):
|
|
engine, table = get_table(db_path, type_=Location)
|
|
datas = engine.execute(table.select()).fetchall()
|
|
yield from (Location(**d) for d in datas)
|
|
|
|
def test(tmp_path):
|
|
tdir = Path(tmp_path)
|
|
tdb = tdir / 'test.sqlite'
|
|
test_limit = 100
|
|
test_src = Path('/L/tmp/LocationHistory.json')
|
|
|
|
# TODO meh, double loading, but for now fine
|
|
with test_src.open('r') as fo:
|
|
real_locs = list(islice(_load_locations(fo), 0, test_limit))
|
|
|
|
cache_locs(source=test_src, db_path=tdb, limit=test_limit)
|
|
cached_locs = list(iter_db_locs(tdb))
|
|
assert len(cached_locs) == test_limit
|
|
assert real_locs == cached_locs
|
|
|
|
from kython.ktyping import PathIsh
|
|
|
|
# TODO what if we want dynamic path??
|
|
# dbcache = make_dbcache('/L/tmp/test.db', hashf=lambda p: p) # TODO FIXME?
|
|
|
|
Hash = str
|
|
|
|
# TODO hash is a bit misleading
|
|
# TODO perhaps another table is the way to go...
|
|
|
|
# TODO careful about concurrent access?
|
|
def read_hash(db_path: Path) -> Optional[Hash]:
|
|
hash_file = db_path.with_suffix('.hash')
|
|
if not hash_file.exists():
|
|
return None
|
|
return hash_file.read_text()
|
|
|
|
# TODO not sure if there is any way to guarantee atomic reading....
|
|
# unless it happens automatically due to unlink logic?
|
|
# TODO need to know entry type?
|
|
# TODO or, we can just encode that in names. that way no need for atomic stuff
|
|
|
|
# TODO give a better name
|
|
class Alala:
|
|
def __init__(self, db_path: Path, type_) -> None:
|
|
self.db = sa.create_engine(f'sqlite:///{db_path}')
|
|
self.engine = self.db.connect() # TODO do I need to tear anything down??
|
|
self.meta = sa.MetaData(self.engine)
|
|
self.table_hash = sa.Table('hash' , self.meta, sa.Column('value', sa.types.String))
|
|
|
|
schema = make_schema(type_)
|
|
self.table_data = sa.Table('table', self.meta, *schema)
|
|
self.meta.create_all()
|
|
|
|
|
|
def get_dbcache_logger():
|
|
return logging.getLogger('dbcache')
|
|
|
|
# TODO ugh. there should be a nicer way to wrap that...
|
|
def make_dbcache(db_path: PathIsh, hashf, type_):
|
|
logger = get_dbcache_logger()
|
|
db_path = Path(db_path)
|
|
def dec(func):
|
|
@functools.wraps(func)
|
|
def wrapper(key):
|
|
# TODO FIXME make sure we have exclusive write lock
|
|
|
|
alala = Alala(db_path, type_)
|
|
engine = alala.engine
|
|
|
|
prev_hashes = engine.execute(alala.table_hash.select()).fetchall()
|
|
if len(prev_hashes) > 1:
|
|
raise RuntimeError(f'Multiple hashes! {prev_hashes}')
|
|
|
|
prev_hash: Optional[Hash]
|
|
if len(prev_hashes) == 0:
|
|
prev_hash = None
|
|
else:
|
|
prev_hash = prev_hashes[0][0] # TODO ugh, returns a tuple...
|
|
logger.debug('previous hash: %s', prev_hash)
|
|
|
|
h = hashf(key)
|
|
logger.debug('current hash: %s', h)
|
|
assert h is not None # just in case
|
|
|
|
with engine.begin() as transaction:
|
|
if h == prev_hash:
|
|
rows = engine.execute(alala.table_data.select()).fetchall()
|
|
return [type_(**row) for row in rows]
|
|
else:
|
|
datas = func(key)
|
|
if len(datas) > 0:
|
|
engine.execute(alala.table_data.insert().values(datas)) # TODO chunks??
|
|
|
|
# TODO FIXME insert and replace instead
|
|
engine.execute(alala.table_hash.delete())
|
|
engine.execute(alala.table_hash.insert().values([{'value': h}]))
|
|
return datas
|
|
return wrapper
|
|
|
|
# TODO FIXME engine is leaking??
|
|
return dec
|
|
|
|
|
|
def hashf(path: Path) -> Hash:
|
|
mt = int(path.stat().st_mtime)
|
|
return f'{path}.{mt}'
|
|
|
|
dbcache = make_dbcache('test.sqlite', hashf=hashf, type_=Location)
|
|
|
|
@dbcache
|
|
def _xxx_locations(path: Path):
|
|
with path.open('r') as fo:
|
|
return list(islice(_load_locations(fo), 0, 100))
|
|
|
|
|
|
def xxx_locations():
|
|
test_src = Path('/L/tmp/LocationHistory.json')
|
|
return _xxx_locations(test_src)
|
|
|
|
|
|
def main():
|
|
from kython import setup_logzero
|
|
setup_logzero(get_logger(), level=logging.DEBUG)
|
|
setup_logzero(get_dbcache_logger(), level=logging.DEBUG)
|
|
|
|
src_path = Path('hi')
|
|
|
|
db_path = Path('test.sqlite')
|
|
# if db_path.exists():
|
|
# db_path.unlink()
|
|
|
|
res = xxx_locations()
|
|
# new_wrapped = dbcache_worker(db_path=db_path, hashf=hashf, type_=Location, wrapped=wrapped)
|
|
# res = new_wrapped(src_path)
|
|
print(res)
|
|
|
|
# cache_locs(source=Path('/L/tmp/LocationHistory.json'), db_path=db_path)
|
|
# locs = iter_db_locs(db_path)
|
|
# print(len(list(locs)))
|
|
|
|
if __name__ == '__main__':
|
|
main()
|