class ProxyTemplate(object):
def __init__(self, proxy_setting):
self._proxy = DictToken(proxy_setting or {})
def render(self, context):
rendered = self._proxy.render(context)
return get_proxy_info(rendered)
class CheckpointManagerAdapter(tacm.TACheckPointMgr):
"""Wrap TACheckPointMgr for custom usage"""
def __init__(self, namespaces, content, meta_config, task_config):
super(CheckpointManagerAdapter, self).__init__(meta_config, task_config)
if isinstance(namespaces, (list, tuple)):
self.namespaces = (_Token(t) for t in namespaces)
else:
self.namespaces = [_Token(namespaces)]
self.content = DictToken(content)
def _namespaces_for(self, ctx):
return [item.render(ctx) for item in self.namespaces]
def save(self, ctx):
"""Save checkpoint"""
super(CheckpointManagerAdapter, self).update_ckpt(
ckpt=self.content.render(ctx),
namespaces=self._namespaces_for(ctx)
)
def load(self, ctx):
"""Load checkpoint"""
namespaces = self._namespaces_for(ctx)
checkpoint = super(CheckpointManagerAdapter, self).get_ckpt(namespaces)
if checkpoint is None:
logger.info('No existing checkpoint found')
checkpoint = {}
return checkpoint
class RequestTemplate(object):
def __init__(self, request):
if not request:
raise ValueError('The request is none')
url = request.get('url')
if not url:
raise ValueError("The request doesn't contain a url or it's empty")
self.url = _Token(url)
self.nextpage_url = _Token(request.get('nextpage_url', url))
self.headers = DictToken(request.get('headers', {}))
# Request body could be string or dict
body = request.get('body')
if isinstance(body, dict):
self.body = DictToken(body)
elif isinstance(body, six.string_types):
self.body = _Token(body)
else:
if body:
logger.warning('Invalid request body: %s', body)
self.body = None
method = request.get('method', 'GET')
if not method or method.upper() not in ('GET', 'POST'):
raise ValueError(
'Unsupported value for request method: {}'.format(method))
self.method = _Token(method)
self.count = 0
def reset(self):
self.count = 0
def render(self, context):
if self.count == 0 or not self.nextpage_url:
url = self.url.render(context)
else:
url = self.nextpage_url.render(context)
self.count += 1
return Request(url=url,
method=self.method.render(context),
headers=self.headers.render(context),
body=self.body.render(context) if self.body else None)