my.fbmessenger.export: use context manager to properly close sqlite connection

This commit is contained in:
Dima Gerasimov 2023-02-08 01:44:20 +00:00 committed by karlicoss
parent bb5ad2b6ac
commit fb0c1289f0

View file

@ -7,10 +7,13 @@ REQUIRES = [
'git+https://github.com/karlicoss/fbmessengerexport', 'git+https://github.com/karlicoss/fbmessengerexport',
] ]
from contextlib import ExitStack, contextmanager
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import Iterator from typing import Iterator
from my.core import PathIsh, Res, stat, Stats
from my.core.warnings import high
from my.config import fbmessenger as user_config from my.config import fbmessenger as user_config
import fbmessengerexport.dal as messenger import fbmessengerexport.dal as messenger
@ -22,7 +25,6 @@ _new_section = getattr(user_config, 'fbmessengerexport', None)
_old_attr = getattr(user_config, 'export_db', None) _old_attr = getattr(user_config, 'export_db', None)
if _new_section is None and _old_attr is not None: if _new_section is None and _old_attr is not None:
from my.core.warnings import high
high("""DEPRECATED! Please modify your fbmessenger config to look like: high("""DEPRECATED! Please modify your fbmessenger config to look like:
class fbmessenger: class fbmessenger:
@ -35,24 +37,26 @@ class fbmessenger:
### ###
from ..core import PathIsh
@dataclass @dataclass
class config(user_config.fbmessengerexport): class config(user_config.fbmessengerexport):
export_db: PathIsh export_db: PathIsh
def _dal() -> messenger.DAL: @contextmanager
return messenger.DAL(config.export_db) def _dal() -> Iterator[messenger.DAL]:
model = messenger.DAL(config.export_db)
with ExitStack() as stack:
if hasattr(model, '__dal__'): # defensive to support legacy fbmessengerexport
stack.enter_context(model)
yield model
from ..core import Res
def messages() -> Iterator[Res[messenger.Message]]: def messages() -> Iterator[Res[messenger.Message]]:
model = _dal() with _dal() as model:
for t in model.iter_threads(): for t in model.iter_threads():
yield from t.iter_messages() yield from t.iter_messages()
from ..core import stat, Stats
def stats() -> Stats: def stats() -> Stats:
return stat(messages) return stat(messages)
@ -75,11 +79,9 @@ def dump_chat_history(where: PathIsh) -> None:
p = Path(where) p = Path(where)
assert not p.exists() or p.is_dir() assert not p.exists() or p.is_dir()
model = _dal()
from shutil import rmtree from shutil import rmtree
from tempfile import TemporaryDirectory from tempfile import TemporaryDirectory
with TemporaryDirectory() as tdir: with TemporaryDirectory() as tdir, _dal() as model:
td = Path(tdir) td = Path(tdir)
_dump_helper(model, td) _dump_helper(model, td)