102 lines
3.1 KiB
Python
102 lines
3.1 KiB
Python
|
from abc import ABCMeta, abstractmethod
|
||
|
|
from six import add_metaclass
|
||
|
|
from functools import partial, wraps
|
||
|
|
|
||
|
|
from prometheus_client import CollectorRegistry, Gauge, Counter, push_to_gateway
|
||
|
|
|
||
|
|
from util.abchelpers import nooper
|
||
|
|
|
||
|
|
|
||
|
|
@add_metaclass(ABCMeta)
|
||
|
|
class ProgressReporter(object):
|
||
|
|
""" Implements an interface for reporting progress with the migrations.
|
||
|
|
"""
|
||
|
|
@abstractmethod
|
||
|
|
def report_version_complete(self, success):
|
||
|
|
""" Called when an entire migration is complete. """
|
||
|
|
|
||
|
|
@abstractmethod
|
||
|
|
def report_step_progress(self):
|
||
|
|
""" Called when a single step in the migration has been completed. """
|
||
|
|
|
||
|
|
|
||
|
|
@nooper
|
||
|
|
class NullReporter(ProgressReporter):
|
||
|
|
""" No-op version of the progress reporter, designed for use when no progress
|
||
|
|
reporting endpoint is provided. """
|
||
|
|
|
||
|
|
|
||
|
|
class PrometheusReporter(ProgressReporter):
|
||
|
|
def __init__(self, prom_pushgateway_addr, prom_job, labels, total_steps_num=None):
|
||
|
|
self._total_steps_num = total_steps_num
|
||
|
|
self._completed_steps = 0.0
|
||
|
|
|
||
|
|
registry = CollectorRegistry()
|
||
|
|
|
||
|
|
self._migration_completion_percent = Gauge(
|
||
|
|
'migration_completion_percent',
|
||
|
|
'Estimate of the completion percentage of the job',
|
||
|
|
registry=registry,
|
||
|
|
)
|
||
|
|
self._migration_complete_total = Counter(
|
||
|
|
'migration_complete_total',
|
||
|
|
'Binary value of whether or not the job is complete',
|
||
|
|
registry=registry,
|
||
|
|
)
|
||
|
|
self._migration_failed_total = Counter(
|
||
|
|
'migration_failed_total',
|
||
|
|
'Binary value of whether or not the job has failed',
|
||
|
|
registry=registry,
|
||
|
|
)
|
||
|
|
self._migration_items_completed_total = Counter(
|
||
|
|
'migration_items_completed_total',
|
||
|
|
'Number of items this migration has completed',
|
||
|
|
registry=registry,
|
||
|
|
)
|
||
|
|
|
||
|
|
self._push = partial(push_to_gateway,
|
||
|
|
prom_pushgateway_addr,
|
||
|
|
job=prom_job,
|
||
|
|
registry=registry,
|
||
|
|
grouping_key=labels,
|
||
|
|
)
|
||
|
|
|
||
|
|
def report_version_complete(self, success=True):
|
||
|
|
if success:
|
||
|
|
self._migration_complete_total.inc()
|
||
|
|
else:
|
||
|
|
self._migration_failed_total.inc()
|
||
|
|
self._migration_completion_percent.set(1.0)
|
||
|
|
|
||
|
|
self._push()
|
||
|
|
|
||
|
|
def report_step_progress(self):
|
||
|
|
self._migration_items_completed_total.inc()
|
||
|
|
|
||
|
|
if self._total_steps_num is not None:
|
||
|
|
self._completed_steps += 1
|
||
|
|
self._migration_completion_percent = self._completed_steps / self._total_steps_num
|
||
|
|
|
||
|
|
self._push()
|
||
|
|
|
||
|
|
|
||
|
|
class ProgressWrapper(object):
|
||
|
|
def __init__(self, delegate_module, progress_monitor):
|
||
|
|
self._delegate_module = delegate_module
|
||
|
|
self._progress_monitor = progress_monitor
|
||
|
|
|
||
|
|
def __getattr__(self, attr_name):
|
||
|
|
# Will raise proper attribute error
|
||
|
|
maybe_callable = self._delegate_module.__dict__[attr_name]
|
||
|
|
if callable(maybe_callable):
|
||
|
|
# Build a callable which when executed places the request
|
||
|
|
# onto a queue
|
||
|
|
@wraps(maybe_callable)
|
||
|
|
def wrapped_method(*args, **kwargs):
|
||
|
|
result = maybe_callable(*args, **kwargs)
|
||
|
|
self._progress_monitor.report_step_progress()
|
||
|
|
return result
|
||
|
|
|
||
|
|
return wrapped_method
|
||
|
|
return maybe_callable
|