quickstart.py

#!/usr/bin/env python

# Copyright 2016 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import argparse


def end_to_end(project_id, topic_name, subscription_name, num_messages):
    # [START pubsub_end_to_end]
    import time

    from google.cloud import pubsub_v1

    # TODO project_id = "Your Google Cloud Project ID"
    # TODO topic_name = "Your Pub/Sub topic name"
    # TODO num_messages = number of messages to test end-to-end

    # Instantiates a publisher and subscriber client
    publisher = pubsub_v1.PublisherClient()
    subscriber = pubsub_v1.SubscriberClient()

    # The `topic_path` method creates a fully qualified identifier
    # in the form `projects/{project_id}/topics/{topic_name}`
    topic_path = subscriber.topic_path(project_id, topic_name)

    # The `subscription_path` method creates a fully qualified identifier
    # in the form `projects/{project_id}/subscriptions/{subscription_name}`
    subscription_path = subscriber.subscription_path(
        project_id, subscription_name)

    # Create the topic.
    topic = publisher.create_topic(topic_path)
    print('\nTopic created: {}'.format(topic.name))

    # Create a subscription.
    subscription = subscriber.create_subscription(
        subscription_path, topic_path)
    print('\nSubscription created: {}\n'.format(subscription.name))

    publish_begin = time.time()

    # Publish messages.
    for n in range(num_messages):
        data = u'Message number {}'.format(n)
        # Data must be a bytestring
        data = data.encode('utf-8')
        # When you publish a message, the client returns a future.
        future = publisher.publish(topic_path, data=data)
        print('Published {} of message ID {}.'.format(data, future.result()))

    publish_time = time.time() - publish_begin

    messages = set()

    def callback(message):
        print('Received message: {}'.format(message))
        # Unacknowledged messages will be sent again.
        message.ack()
        messages.add(message)

    subscribe_begin = time.time()

    # Receive messages. The subscriber is nonblocking.
    subscriber.subscribe(subscription_path, callback=callback)

    print('\nListening for messages on {}...\n'.format(subscription_path))

    while True:
        if len(messages) == num_messages:
            subscribe_time = time.time() - subscribe_begin
            print("\nReceived all messages.")
            print("Publish time lapsed: {:.2f}s.".format(publish_time))
            print("Subscribe time lapsed: {:.2f}s.".format(subscribe_time))
            break
        else:
            # Sleeps the thread at 50Hz to save on resources.
            time.sleep(1. / 50)
    # [END pubsub_end_to_end]


if __name__ == '__main__':

    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument('project_id', help='Your Google Cloud project ID')
    parser.add_argument('topic_name', help='Your topic name')
    parser.add_argument('subscription_name', help='Your subscription name')
    parser.add_argument('num_msgs', type=int, help='Number of test messages')

    args = parser.parse_args()

    end_to_end(args.project_id, args.topic_name, args.subscription_name,
               args.num_msgs)