demoweb/venv/lib/python3.7/site-packages/pubcontrol/pubcontrolclient.py

#    pubcontrolclient.py
#    ~~~~~~~~~
#    This module implements the PubControlClient class.
#    :authors: Justin Karneges, Konstantin Bokarius.
#    :copyright: (c) 2015 by Fanout, Inc.
#    :license: MIT, see LICENSE for more details.

import sys
import json
import copy
from base64 import b64encode
import threading
from collections import deque
import requests
from .pubsubmonitor import PubSubMonitor
from .utilities import _gen_auth_jwt_header

try:
	import ndg.httpsclient
except ImportError:
	ndg = None

# TODO: Remove in the future when most users have Python >= 2.7.9.
try:
	import urllib3
except ImportError:
	try:
		from requests.packages import urllib3
	except ImportError:
		pass
try:
	urllib3.disable_warnings()
except NameError:
	pass
except AttributeError:
	pass

# The PubControlClient class allows consumers to publish either synchronously
# or asynchronously to an endpoint of their choice. The consumer wraps a Format
# class instance in an Item class instance and passes that to the publish
# method. The async publish method has an optional callback parameter that
# is called after the publishing is complete to notify the consumer of the
# result. Optionally provide JWT authentication claim and key information.
# If require_subscribers is set to True then channel subscription monitoring
# will be enabled and only channels that are subscribed to will be published
# to.
class PubControlClient(object):

	# Initialize this class with a URL representing the publishing endpoint.
	def __init__(self, uri, auth_jwt_claim=None,
			auth_jwt_key=None, require_subscribers=False, sub_callback=None):
		self.uri = uri
		self.lock = threading.Lock()
		self.thread = None
		self.thread_cond = None
		self.req_queue = deque()
		self.auth_basic_user = None
		self.auth_basic_pass = None
		self.auth_jwt_claim = auth_jwt_claim
		self.auth_jwt_key = auth_jwt_key
		self.requests_session = requests.session()
		self.sub_monitor = None
		self.closed = False
		if require_subscribers:
			self.sub_monitor = PubSubMonitor(uri, auth_jwt_claim, auth_jwt_key, sub_callback)

	# Call this method and pass a username and password to use basic
	# authentication with the configured endpoint.
	def set_auth_basic(self, username, password):
		self._verify_notclosed()
		self.lock.acquire()
		self.auth_basic_user = username
		self.auth_basic_pass = password
		self.lock.release()

	# Call this method and pass a claim and key to use JWT authentication
	# with the configured endpoint.
	def set_auth_jwt(self, claim, key):
		self._verify_notclosed()
		self.lock.acquire()
		self.auth_jwt_claim = claim
		self.auth_jwt_key = key
		self.lock.release()

	# The publish method for publishing the specified item to the specified
	# channel on the configured endpoint. The blocking parameter indicates
	# whether the the callback method should be blocking or non-blocking. The
	# callback parameter is optional and will be passed the publishing results
	# after publishing is complete. Note that the callback executes on a
	# separate thread. If require_subscribers was set to True then the message
	# will only be published if the channel is subscribed to. If the sub_monitor
	# instance failed to retrieve subscriber information then an error will be
	# raised.
	def publish(self, channel, item, blocking=False, callback=None):
		self._verify_notclosed()
		if self.sub_monitor and self.sub_monitor.is_closed():
			if callback:
				callback(False, 'failed to retrieve channel subscribers')
			else:
				raise ValueError('failed to retrieve channel subscribers')
		elif self.sub_monitor and not self.sub_monitor.is_channel_subscribed_to(channel):
			if callback:
				callback(True, '')
			return
		i = item.export()
		i['channel'] = channel
		if blocking:
			self.lock.acquire()
			uri = self.uri
			auth = self._gen_auth_header()
			self.lock.release()
			self._pubcall(uri, auth, [i])
		else:
			self.lock.acquire()
			uri = self.uri
			auth = self._gen_auth_header()
			self._ensure_thread()
			self.lock.release()
			self._queue_req(('pub', uri, auth, i, callback))

	# This method is a blocking method that ensures that all asynchronous
	# publishing is complete prior to returning and allowing the consumer to
	# proceed.
	def wait_all_sent(self):
		self.lock.acquire()
		if self.thread is not None:
			self._queue_req(('stop',))
			self.thread.join()
			self.thread = None
		self.lock.release()

	# DEPRECATED: The finish method is now deprecated in favor of the more
	# descriptive wait_all_sent() method.
	def finish(self):
		self.wait_all_sent()

	# This method is meant to close the PubControlClient instance and should
	# implemented when needed for future features. Currently this method
	# is simply a passthrough to wait_all_sent().
	def close(self):
		self.lock.acquire()
		self.closed = True
		if self.sub_monitor:
			self.sub_monitor.close()
		self.lock.release()
		self.wait_all_sent()

	# This method makes an HTTP request to an endpoint relative to the base
	# URI, using configured authentication. Returns a tuple of
	# (status code, headers, body).
	def http_call(self, endpoint, data, headers={}):
		uri = self.uri + endpoint

		send_headers = copy.deepcopy(headers)

		auth_header = self._gen_auth_header()
		if auth_header:
			send_headers['Authorization'] = auth_header

		try:
			return self._make_http_request(uri, data, send_headers)
		except Exception as e:
			raise ValueError('failed during http call: ' + str(e))

	# An internal method for verifying that the PubControlClient instance
	# has not been closed via the close() method. If it has then an error
	# is raised.
	def _verify_notclosed(self):
		if self.closed:
			raise ValueError('pubcontrolclient instance is closed')

	# An internal method used to generate an authorization header. The
	# authorization header is generated based on whether basic or JWT
	# authorization information was provided via the publicly accessible
	# 'set_*_auth' methods defined above.
	def _gen_auth_header(self):
		if self.auth_basic_user:
			return 'Basic ' + str(b64encode(('%s:%s' % (self.auth_basic_user, self.auth_basic_pass)).encode('ascii')))
		elif self.auth_jwt_claim:
			return _gen_auth_jwt_header(self.auth_jwt_claim, self.auth_jwt_key)
		else:
			return None

	# An internal method that ensures that asynchronous publish calls are
	# properly processed. This method initializes the required class fields,
	# starts the pubworker worker thread, and is meant to execute only when
	# the consumer makes an asynchronous publish call.
	def _ensure_thread(self):
		if self.thread is None:
			self.thread_cond = threading.Condition()
			self.thread = threading.Thread(target=self._pubworker)
			self.thread.daemon = True
			self.thread.start()

	# An internal method for adding an asynchronous publish request to the
	# publishing queue. This method will also activate the pubworker worker
	# thread to make sure that it process any and all requests added to
	# the queue.
	def _queue_req(self, req):
		self.thread_cond.acquire()
		self.req_queue.append(req)
		self.thread_cond.notify()
		self.thread_cond.release()

	# An internal method for preparing the HTTP POST request for publishing
	# data to the endpoint. This method accepts the URI endpoint, authorization
	# header, and a list of items to publish.
	def _pubcall(self, uri, auth_header, items):
		uri = uri + '/publish/'

		headers = dict()
		if auth_header:
			headers['Authorization'] = auth_header
		headers['Content-Type'] = 'application/json'

		content = dict()
		content['items'] = items
		content_raw = json.dumps(content)

		try:
			if isinstance(content_raw, unicode):
				content_raw = content_raw.encode('utf-8')
		except NameError:
			if isinstance(content_raw, str):
				content_raw = content_raw.encode('utf-8')

		try:
			self._make_http_request(uri, content_raw, headers)
		except Exception as e:
			raise ValueError('failed to publish: ' + str(e))

	# An internal method for making an HTTP request to the specified URI
	# with the specified content and headers.
	def _make_http_request(self, uri, data, headers):
		if sys.version_info >= (2, 7, 9) or (ndg and ndg.httpsclient):
			res = self.requests_session.post(uri, headers=headers, data=data)
		else:
			res = self.requests_session.post(uri, headers=headers, data=data, verify=False)
		self._verify_status_code(res.status_code, res.text)
		return (res.status_code, res.headers, res.text)

	# An internal method for ensuring a successful status code is returned
	# from the server.
	def _verify_status_code(self, code, message):
		if code < 200 or code >= 300:
			raise ValueError('received failed status code ' + str(code) +
					' with message: ' + message)

	# An internal method for publishing a batch of requests. The requests are
	# parsed for the URI, authorization header, and each request is published
	# to the endpoint. After all publishing is complete, each callback
	# corresponding to each request is called (if a callback was originally
	# provided for that request) and passed a result indicating whether that
	# request was successfully published.
	def _pubbatch(self, reqs):
		assert(len(reqs) > 0)
		uri = reqs[0][0]
		auth_header = reqs[0][1]
		items = list()
		callbacks = list()
		for req in reqs:
			items.append(req[2])
			callbacks.append(req[3])

		try:
			self._pubcall(uri, auth_header, items)
			result = (True, '')
		except Exception as e:
			try:
				result = (False, e.message)
			except AttributeError:
				result = (False, str(e))

		for c in callbacks:
			if c:
				c(result[0], result[1])

	# An internal method that is meant to run as a separate thread and process
	# asynchronous publishing requests. The method runs continously and
	# publishes requests in batches containing a maximum of 10 requests. The
	# method completes and the thread is terminated only when a 'stop' command
	# is provided in the request queue.
	def _pubworker(self):
		quit = False
		while not quit:
			self.thread_cond.acquire()

			# if no requests ready, wait for one
			if len(self.req_queue) == 0:
				self.thread_cond.wait()

				# still no requests after notification? start over
				if len(self.req_queue) == 0:
					self.thread_cond.release()
					continue

			reqs = list()
			while len(self.req_queue) > 0 and len(reqs) < 10:
				m = self.req_queue.popleft()
				if m[0] == 'stop':
					quit = True
					break
				reqs.append((m[1], m[2], m[3], m[4]))

			self.thread_cond.release()

			if len(reqs) > 0:
				self._pubbatch(reqs)