def test_get_pubsub_client_with_service_account(self):
"""Tests the client obtained by the service account method."""
client = get_pubsub_client()
project = ('projects/%s' %
os.environ.get(TEST_PROJECT_ENV, DEFAULT_TEST_PROJECT))
client.projects().topics().list(project=project).execute(
num_retries=self.RETRY)
# Providing an Http object this time.
self.assertIsNotNone(get_pubsub_client(http=httplib2.Http()))
def send_loop(client, q, topic, retry, logger, format_func,
publish_body): # pragma: NO COVER
"""Process loop for indefinitely sending logs to Cloud Pub/Sub.
Args:
client: Pub/Sub client. If it's None, a new client will be created.
q: mp.JoinableQueue instance to get the message from.
topic: Cloud Pub/Sub topic name to send the logs.
retry: How many times to retry upon Cloud Pub/Sub API failure.
logger: A logger for informing failures within this function.
format_func: A callable for formatting the logs.
publish_body: A callable for sending the logs.
"""
if client is None:
client = get_pubsub_client()
while True:
try:
logs = q.get()
except Empty:
continue
try:
body = {'messages': [{'data': format_func(r)} for r in logs]}
publish_body(client, body, topic)
except errors.RecoverableError as e:
# Records the exception and puts the logs back to the deque
# and prints the exception to stderr.
q.put(logs)
logger.exception(e)
except Exception as e:
logger.exception(e)
logger.warn('There was a non recoverable error, exiting.')
return
q.task_done()
def setUp(self):
self.client = get_pubsub_client()
self.project = os.environ.get(TEST_PROJECT_ENV, DEFAULT_TEST_PROJECT)
self.topic = 'projects/%s/topics/test-topic-%f' % (self.project,
time.time())
self.nonexistence = 'projects/%s/topics/nonexistence' % self.project
try:
self.client.projects().topics().create(name=self.topic,
body={}).execute()
except errors.HttpError as e:
if e.resp.status == 409:
pass
else:
raise
def __init__(self,
project_id,
topic,
capacity=DEFAULT_BATCH_NUM,
retry=DEFAULT_RETRY_COUNT,
flush_level=logging.CRITICAL,
buf_hard_limit=-1,
client=None,
publish_body=publish_body):
"""The constructor of the handler.
Args:
topic: Cloud Pub/Sub topic name to send the logs.
capacity: The maximum buffer size, defaults to 1000.
retry: How many times to retry upon Cloud Pub/Sub API failure,
defaults to 5.
flush_level: Minimum log level that, when seen, will flush the logs,
defaults to logging.CRITICAL.
buf_hard_limit: Maximum buffer size to hold the logs, defaults to -1
which means unlimited size.
client: An optional Cloud Pub/Sub client to use. If not set, one is
built automatically, defaults to None.
publish_body: A callable for publishing the Pub/Sub message,
just for testing and benchmarking purposes.
"""
super(PubsubHandler, self).__init__(capacity)
self._topic = topic
self._project_id = project_id
self._retry = retry
self._flush_level = flush_level
self._buf_hard_limit = buf_hard_limit
self._publish_body = publish_body
if client:
self._client = client
else:
self._client = get_pubsub_client()
if not check_topic(self._client, project_id, topic):
raise EnvironmentError(
'Failed to confirm the existence of the topic "%s".' % topic)
def __init__(self,
topic,
worker_num=DEFAULT_POOL_SIZE,
retry=DEFAULT_RETRY_COUNT,
client=None,
publish_body=publish_body,
stderr_logger=None):
"""The constructor of the handler.
Args:
topic: Cloud Pub/Sub topic name to send the logs.
worker_num: The number of workers, defaults to 1.
retry: How many times to retry upon Cloud Pub/Sub API failure,
defaults to 5.
client: An optional Cloud Pub/Sub client to use. If not set, one is
built automatically, defaults to None.
publish_body: A callable for publishing the Pub/Sub message,
just for testing and benchmarking purposes.
stderr_logger: A logger for informing failures with this
logger, defaults to None and if not specified, a last
resort logger will be used.
"""
super(AsyncPubsubHandler, self).__init__()
self._q = mp.JoinableQueue()
self._batch_size = BATCH_SIZE
self._buf = []
if not check_topic(client or get_pubsub_client(), topic, retry):
raise EnvironmentError(
'Failed to confirm the existence of the topic "%s".' % topic)
if not stderr_logger:
stderr_logger = logging.Logger('last_resort')
stderr_logger.addHandler(logging.StreamHandler())
for _ in range(worker_num):
p = mp.Process(target=send_loop,
args=(client, self._q, topic, retry, stderr_logger,
self.format, publish_body))
p.daemon = True
p.start()
def test_get_pubsub_client_with_scoped_credentials(self):
"""Tests the client obtained by scoped credentials."""
credentials = GoogleCredentials.get_application_default()
credentials = credentials.create_scoped(PUBSUB_SCOPES)
self.assertIsNotNone(get_pubsub_client(credentials=credentials))