def create_occurrence_subscription(subscription_id, project_id):
"""Creates a new Pub/Sub subscription object listening to the
Container Analysis Occurrences topic."""
# subscription_id := 'my-occurrences-subscription'
# project_id = 'my-gcp-project'
from google.api_core.exceptions import AlreadyExists
from google.cloud.pubsub import SubscriberClient
topic_id = 'container-analysis-occurrences-v1'
client = SubscriberClient()
topic_name = f"projects/{project_id}/topics/{topic_id}"
subscription_name = client.subscription_path(project_id, subscription_id)
success = True
try:
client.create_subscription({
"name": subscription_name,
"topic": topic_name
})
except AlreadyExists:
# if subscription already exists, do nothing
pass
else:
success = False
return success
def _create_client(service_account_file=None, service_account_json=None):
if service_account_file is not None:
return SubscriberClient.from_service_account_file(
service_account_file)
if service_account_json is not None:
return SubscriberClient.from_service_account_json(
service_account_json)
return SubscriberClient()
def subscribe_with_avro_schema(project_id,
subscription_id,
avsc_file,
timeout=None):
"""Receive and decode messages sent to a topic with an Avro schema."""
# [START pubsub_subscribe_avro_records]
import avro
from avro.io import BinaryDecoder, DatumReader
from concurrent.futures import TimeoutError
import io
import json
from google.cloud.pubsub import SubscriberClient
# TODO(developer)
# project_id = "your-project-id"
# subscription_id = "your-subscription-id"
# avsc_file = "path/to/an/avro/schema/file/(.avsc)/formatted/in/json"
# Number of seconds the subscriber listens for messages
# timeout = 5.0
subscriber = SubscriberClient()
subscription_path = subscriber.subscription_path(project_id,
subscription_id)
avro_schema = avro.schema.parse(open(avsc_file, "rb").read())
def callback(message):
# Get the message serialization type.
encoding = message.attributes.get("googclient_schemaencoding")
# Deserialize the message data accordingly.
if encoding == "BINARY":
bout = io.BytesIO(message.data)
decoder = BinaryDecoder(bout)
reader = DatumReader(avro_schema)
message_data = reader.read(decoder)
print(f"Received a binary-encoded message:\n{message_data}")
elif encoding == "JSON":
message_data = json.loads(message.data)
print(f"Received a JSON-encoded message:\n{message_data}")
else:
print(f"Received a message with no encoding:\n{message}")
message.ack()
streaming_pull_future = subscriber.subscribe(subscription_path,
callback=callback)
print(f"Listening for messages on {subscription_path}..\n")
# Wrap subscriber in a 'with' block to automatically call close() when done.
with subscriber:
try:
# When `timeout` is not set, result() will block indefinitely,
# unless an exception occurs first.
streaming_pull_future.result(timeout=timeout)
except TimeoutError:
streaming_pull_future.cancel()
def subscribe_with_proto_schema(
project_id: str, subscription_id: str, timeout: float
) -> None:
"""Receive and decode messages sent to a topic with a protobuf schema."""
# [[START pubsub_subscribe_proto_messages]
from concurrent.futures import TimeoutError
from google.cloud.pubsub import SubscriberClient
from google.protobuf.json_format import Parse
from utilities import us_states_pb2
# TODO(developer)
# project_id = "your-project-id"
# subscription_id = "your-subscription-id"
# Number of seconds the subscriber listens for messages
# timeout = 5.0
subscriber = SubscriberClient()
subscription_path = subscriber.subscription_path(project_id, subscription_id)
# Instantiate a protoc-generated class defined in `us-states.proto`.
state = us_states_pb2.StateProto()
def callback(message: pubsub_v1.subscriber.message.Message) -> None:
# Get the message serialization type.
encoding = message.attributes.get("googclient_schemaencoding")
# Deserialize the message data accordingly.
if encoding == "BINARY":
state.ParseFromString(message.data)
print("Received a binary-encoded message:\n{state}")
elif encoding == "JSON":
Parse(message.data, state)
print(f"Received a JSON-encoded message:\n{state}")
else:
print(f"Received a message with no encoding:\n{message}")
message.ack()
streaming_pull_future = subscriber.subscribe(subscription_path, callback=callback)
print(f"Listening for messages on {subscription_path}..\n")
# Wrap subscriber in a 'with' block to automatically call close() when done.
with subscriber:
try:
# When `timeout` is not set, result() will block indefinitely,
# unless an exception occurs first.
streaming_pull_future.result(timeout=timeout)
except TimeoutError:
streaming_pull_future.cancel() # Trigger the shutdown.
streaming_pull_future.result() # Block until the shutdown is complete.
def pubsub(subscription_id, timeout_seconds, project_id):
"""Respond to incoming occurrences using a Cloud Pub/Sub subscription."""
# subscription_id := 'my-occurrences-subscription'
# timeout_seconds = 20
# project_id = 'my-gcp-project'
import time
from google.cloud.pubsub import SubscriberClient
client = SubscriberClient()
subscription_name = client.subscription_path(project_id, subscription_id)
receiver = MessageReceiver()
client.subscribe(subscription_name, receiver.pubsub_callback)
# listen for 'timeout' seconds
for _ in range(timeout_seconds):
time.sleep(1)
# print and return the number of pubsub messages received
print(receiver.msg_count)
return receiver.msg_count
def pull(self, max_messages=0):
from google.cloud.pubsub import SubscriberClient
subscriber = SubscriberClient()
sub_path = subscriber.subscription_path(self.project, self.full_name)
if max_messages == 0:
response = subscriber.pull(sub_path)
else:
response = subscriber.pull(sub_path, max_messages=max_messages)
ack_ids = [msg.ack_id for msg in response.received_messages]
if ack_ids:
subscriber.acknowledge(sub_path, ack_ids)
return [m.message.data for m in response.received_messages]
def test_pubsub(self):
# create topic if needed
client = SubscriberClient()
try:
topic_id = 'container-analysis-occurrences-v1'
topic_name = client.topic_path(PROJECT_ID, topic_id)
publisher = PublisherClient()
publisher.create_topic(topic_name)
except AlreadyExists:
pass
subscription_id = 'container-analysis-test-{}'.format(uuid.uuid4())
subscription_name = client.subscription_path(PROJECT_ID,
subscription_id)
samples.create_occurrence_subscription(subscription_id, PROJECT_ID)
# I can not make it pass with multiple messages. My guess is
# the server started to dedup?
message_count = 1
try:
job_done = threading.Event()
receiver = MessageReceiver(message_count, job_done)
client.subscribe(subscription_name, receiver.pubsub_callback)
for i in range(message_count):
occ = samples.create_occurrence(self.image_url, self.note_id,
PROJECT_ID, PROJECT_ID)
time.sleep(SLEEP_TIME)
samples.delete_occurrence(basename(occ.name), PROJECT_ID)
time.sleep(SLEEP_TIME)
# We saw occational failure with 60 seconds timeout, so we bumped it
# to 180 seconds.
# See also: python-docs-samples/issues/2894
job_done.wait(timeout=180)
print('done. msg_count = {}'.format(receiver.msg_count))
assert message_count <= receiver.msg_count
finally:
# clean up
client.delete_subscription(subscription_name)
def clients():
cleanup()
run_cmd = shlex.split(
f"docker run -d -p 8085:8085 --name {CONTAINER} google/cloud-sdk "
"gcloud beta emulators pubsub start "
f"--project={PROJECT} --host-port=0.0.0.0:8085"
)
subprocess.check_call(run_cmd)
def predicate():
cmd = shlex.split(f"docker logs {CONTAINER}")
logs = subprocess.check_output(cmd, stderr=subprocess.STDOUT)
return b"INFO: Server started, listening on 8085" in logs
wait_for(predicate, 10, period=0.1)
os.environ["PUBSUB_EMULATOR_HOST"] = "localhost:8085"
os.environ["PUBSUB_PROJECT_ID"] = PROJECT
try:
yield (PublisherClient(), SubscriberClient())
finally:
cleanup(fail=True)
def test_pubsub(self):
# create topic if needed
client = SubscriberClient()
try:
topic_id = 'container-analysis-occurrences-v1'
topic_name = client.topic_path(PROJECT_ID, topic_id)
publisher = PublisherClient()
publisher.create_topic(topic_name)
except AlreadyExists:
pass
subscription_id = 'drydockOccurrences'
subscription_name = client.subscription_path(PROJECT_ID,
subscription_id)
samples.create_occurrence_subscription(subscription_id, PROJECT_ID)
tries = 0
success = False
while not success and tries < TRY_LIMIT:
print(tries)
tries += 1
receiver = samples.MessageReceiver()
client.subscribe(subscription_name, receiver.pubsub_callback)
# test adding 3 more occurrences
total_created = 3
for _ in range(total_created):
occ = samples.create_occurrence(self.image_url, self.note_id,
PROJECT_ID, PROJECT_ID)
sleep(SLEEP_TIME)
samples.delete_occurrence(basename(occ.name), PROJECT_ID)
sleep(SLEEP_TIME)
print('done. msg_count = {}'.format(receiver.msg_count))
success = receiver.msg_count == total_created
if receiver.msg_count != total_created:
raise AssertionError
# clean up
client.delete_subscription(subscription_name)
def subscriber_client():
subscriber_client = SubscriberClient()
yield subscriber_client
subscriber_client.close()
def subscriber_client() -> Generator[pubsub_v1.SubscriberClient, None, None]:
subscriber_client = SubscriberClient()
yield subscriber_client
subscriber_client.close()
