support new ReplicationServer as contextmanager

This commit is contained in:
Sarah Hoffmann 2022-04-07 17:58:04 +02:00
parent f78ae969e9
commit 126cabacb8

View File

@ -7,6 +7,7 @@
""" """
Functions for updating a database from a replication source. Functions for updating a database from a replication source.
""" """
from contextlib import contextmanager
import datetime as dt import datetime as dt
from enum import Enum from enum import Enum
import logging import logging
@ -109,8 +110,7 @@ def update(conn, options):
options['import_file'].unlink() options['import_file'].unlink()
# Read updates into file. # Read updates into file.
repl = ReplicationServer(options['base_url']) with _make_replication_server(options['base_url']) as repl:
outhandler = WriteHandler(str(options['import_file'])) outhandler = WriteHandler(str(options['import_file']))
endseq = repl.apply_diffs(outhandler, startseq + 1, endseq = repl.apply_diffs(outhandler, startseq + 1,
max_size=options['max_diff_size'] * 1024) max_size=options['max_diff_size'] * 1024)
@ -130,3 +130,19 @@ def update(conn, options):
seq=endseq, indexed=False) seq=endseq, indexed=False)
return UpdateState.UP_TO_DATE return UpdateState.UP_TO_DATE
def _make_replication_server(url):
""" Returns a ReplicationServer in form of a context manager.
Creates a light wrapper around older versions of pyosmium that did
not support the context manager interface.
"""
if hasattr(ReplicationServer, '__enter__'):
return ReplicationServer(url)
@contextmanager
def get_cm():
yield ReplicationServer(url)
return get_cm()