class GooglePubSubPublisher:
def __init__(self, **kwargs):
super().__init__(**kwargs)
self._logger = logging.getLogger(__name__)
self._publisher = PublisherClient()
def send_message(self, message):
serializer = message.get_serializer()
json_data = JSONRenderer().render(serializer.data)
topic_path = self._publisher.topic_path(
settings.GCLOUD_PUBSUB_PROJECT_ID, serializer.data['type'])
future = self._publisher.publish(topic_path, data=json_data)
future.add_done_callback(self._publish_callback)
def _publish_callback(self, message_future):
if message_future.exception(timeout=30):
self._logger.info(
f'Publishing message on {message_future.exception()} threw an Exception.'
)
else:
self._logger.info(message_future.result())
def create_topic(self, topic: str):
topic_path = self._publisher.topic_path(
settings.GCLOUD_PUBSUB_PROJECT_ID, topic)
try:
self._publisher.create_topic(topic_path)
except AlreadyExists as e:
self._logger.info(f'Topic {topic} already exists.')
class PubSubPublisher(Publisher):
def __init__(self, project_id, credentials_file=None):
logger.info('connecting to pubsub')
if credentials_file:
service_account_info = json.load(open(credentials_file))
publisher_audience = 'https://pubsub.googleapis.com/google.pubsub.v1.Publisher'
credentials_pub = jwt.Credentials.from_service_account_info(
service_account_info, audience=publisher_audience,
)
self._client = PublisherClient(credentials=credentials_pub)
else:
self._client = PublisherClient()
self._project_id = project_id
def _publish(self, topic_id, message, **kwargs):
logger.info('publishing message', topic_id=topic_id)
topic_path = self._client.topic_path(self._project_id, topic_id) # pylint: disable=no-member
response: Future = self._client.publish(topic_path, message.encode('utf-8'), **kwargs)
return response
def publish(self, topic_id, message: bytes, fulfilment_request_transaction_id: str):
response = self._publish(topic_id, message, tx_id=fulfilment_request_transaction_id)
try:
# Resolve the future
message_id = response.result()
logger.info( # pragma: no cover
'message published successfully',
topic_id=topic_id,
message_id=message_id,
fulfilment_request_transaction_id=fulfilment_request_transaction_id,
)
except Exception as ex: # pylint:disable=broad-except
logger.exception(
'message publication failed',
topic_id=topic_id,
)
raise PublicationFailed(ex)
def create_topic(self, topic_id):
try:
logger.info('creating topic')
topic_path = self._client.topic_path(self._project_id, topic_id) # pylint: disable=no-member
self._client.create_topic(request={'name': topic_path}) # pylint: disable=no-member
except AlreadyExists:
logger.info('Topic already exists')
except Exception as ex: # pylint:disable=broad-except
logger.error(
'failed', exc_info=ex,
)
class PubSubPublisher(Publisher):
def __init__(self):
self._client = PublisherClient()
_, self._project_id = google.auth.default()
def _publish(self, topic_id, message):
logger.info("publishing message", topic_id=topic_id)
topic_path = self._client.topic_path(self._project_id, topic_id)
response: Future = self._client.publish(topic_path, message)
return response
def publish(self, topic_id, message: bytes):
response = self._publish(topic_id, message)
try:
# Resolve the future
message_id = response.result()
logger.info( # pragma: no cover
"message published successfully",
topic_id=topic_id,
message_id=message_id,
)
except Exception as ex: # pylint:disable=broad-except
logger.exception(
"message publication failed",
topic_id=topic_id,
)
raise PublicationFailed(ex)
def build_producers(self):
producers = []
project = "testing"
try:
# Create pubsub topics
publisher = PublisherClient()
publisher.create_topic(publisher.topic_path(project, "test10"))
publisher.create_topic(publisher.topic_path(project, "test20"))
# Create pubsub subscriptions
subscriber = SubscriberClient()
subscriber.create_subscription(
subscriber.subscription_path(project, "test10"),
subscriber.topic_path(project, "test10"))
subscriber.create_subscription(
subscriber.subscription_path(project, "test20"),
subscriber.topic_path(project, "test20"))
except AlreadyExists as _:
pass
producers.append(
StreamBuilderFactory.get_producer_pubsub().set_project_id(
project).set_output_topic("test10").build())
producers.append(
StreamBuilderFactory.get_producer_kafka().set_broker_servers(
"localhost:9092").set_output_topic("test1").build())
producers.append(
StreamBuilderFactory.get_producer_kafka().set_broker_servers(
"localhost:9092").set_output_topic("test3").build())
return producers
def ensure_topic_and_subscription():
from google.cloud.pubsub_v1 import PublisherClient, SubscriberClient
publisher_client = PublisherClient()
try:
publisher_client.create_topic(
publisher_client.topic_path(PROJECT_ID, TOPIC_NAME))
except AlreadyExists:
pass
subscriber_client = SubscriberClient()
try:
subscriber_client.create_subscription(
subscriber_client.subscription_path(PROJECT_ID, SUBSCRIPTION_NAME),
publisher_client.topic_path(PROJECT_ID, TOPIC_NAME))
except AlreadyExists:
pass
def delete_topic(
self,
topic: str,
project_id: Optional[str] = None,
fail_if_not_exists: bool = False,
retry: Optional[Retry] = None,
timeout: Optional[float] = None,
metadata: Optional[Sequence[Tuple[str, str]]] = None,
) -> None:
"""
Deletes a Pub/Sub topic if it exists.
:param topic: the Pub/Sub topic name to delete; do not
include the ``projects/{project}/topics/`` prefix.
:type topic: str
:param project_id: Optional, the GCP project ID in which to delete the topic.
If set to None or missing, the default project_id from the GCP connection is used.
:type project_id: str
:param fail_if_not_exists: if set, raise an exception if the topic
does not exist
:type fail_if_not_exists: bool
:param retry: (Optional) A retry object used to retry requests.
If None is specified, requests will not be retried.
:type retry: google.api_core.retry.Retry
:param timeout: (Optional) The amount of time, in seconds, to wait for the request
to complete. Note that if retry is specified, the timeout applies to each
individual attempt.
:type timeout: float
:param metadata: (Optional) Additional metadata that is provided to the method.
:type metadata: Sequence[Tuple[str, str]]]
"""
if not project_id:
raise ValueError("Project ID should be set.")
publisher = self.get_conn()
topic_path = PublisherClient.topic_path(project_id, topic) # pylint: disable=no-member
self.log.info("Deleting topic (path) %s", topic_path)
try:
# pylint: disable=no-member
publisher.delete_topic(
topic=topic_path,
retry=retry,
timeout=timeout,
metadata=metadata,
)
except NotFound:
self.log.warning('Topic does not exist: %s', topic_path)
if fail_if_not_exists:
raise PubSubException(
'Topic does not exist: {}'.format(topic_path))
except GoogleAPICallError as e:
raise PubSubException('Error deleting topic {}'.format(topic), e)
self.log.info("Deleted topic (path) %s", topic_path)
def topic_path(
publisher_client: pubsub_v1.PublisherClient
) -> Generator[str, None, None]:
topic_path = publisher_client.topic_path(PROJECT_ID, TOPIC_ID)
try:
topic = publisher_client.create_topic(request={"name": topic_path})
yield topic.name
except AlreadyExists:
yield topic_path
publisher_client.delete_topic(request={"topic": topic_path})
class PubSubPublisher:
_client_config = {
"interfaces": {
"google.pubsub.v1.Publisher": {
"retry_params": {
"messaging": {
"total_timeout_millis": consts.PUBSUB_TIMEOUT_MS,
}
}
}
}
}
def __init__(self, project_id):
self.project_id = project_id
self.client = PublisherClient(client_config=self._client_config)
self.futures = dict()
def _get_callback(self, f, data):
def callback(f):
if f.exception():
logging.error(f"Please handle {f.exeption()} for {data}.")
if data in self.futures:
self.futures.pop(data)
return callback
def publish_data(self, topic_name, data, timeout=60):
""" Publish Pub/Sub Data
:param topic_name: String name of topic
:param data: String data being processed
:param wait_for_ack: Bool if user wants to wait for ack
:param timeout: Int seconds to wait
"""
if isinstance(data, dict):
data = json.dumps(data)
self.futures.update({data: None})
# When you publish a message, the client returns a future.
topic_path = self.client.topic_path(self.project_id, topic_name)
future = self.client.publish(
topic_path, data=data.encode("utf-8")
)
self.futures[data] = future
# Publish failures shall be handled in the callback function.
future.add_done_callback(self._get_callback(future, data))
return data
def test_create(publisher_client: pubsub_v1.PublisherClient,
capsys: CaptureFixture) -> None:
# The scope of `topic_path` is limited to this function.
topic_path = publisher_client.topic_path(PROJECT_ID, TOPIC_ID)
try:
publisher_client.delete_topic(request={"topic": topic_path})
except NotFound:
pass
publisher.create_topic(PROJECT_ID, TOPIC_ID)
out, _ = capsys.readouterr()
assert f"Created topic: {topic_path}" in out
def topic_path(
publisher_client: pubsub_v1.PublisherClient,
) -> Generator[str, None, None]:
topic_path = publisher_client.topic_path(PROJECT_ID, TOPIC_ID)
try:
topic = publisher_client.get_topic(request={"topic": topic_path})
except NotFound:
topic = publisher_client.create_topic(request={"name": topic_path})
yield topic.name
try:
publisher_client.delete_topic(request={"topic": topic.name})
except NotFound:
pass
def publish_message(response_text, patient_id, event_attributes):
publisher = PublisherClient()
# pylint: disable=no-member
topic_path = publisher.topic_path(
os.environ.get("PROJECT_NAME"), os.environ.get("TOPIC_NAME")
)
data = response_text.encode("utf-8")
future = publisher.publish(
topic_path,
data=data,
**event_attributes,
patient_id=patient_id,
)
message_id = future.result()
logger.info(f"Successfully published message: {message_id}")
print(f"Successfully published message: {message_id}")
def tearDownClass(cls):
super().tearDownClass()
publisher = PublisherClient()
subscriber = SubscriberClient()
for topic in ["test0", "test1"]:
try:
publisher.delete_topic(publisher.topic_path(
cls.project, topic))
except Exception as ex:
raise ex
try:
subscriber.delete_subscription(
subscriber.subscription_path(cls.project, topic))
except Exception as ex:
raise ex
def proto_topic(publisher_client: pubsub_v1.PublisherClient,
proto_schema: str) -> Generator[str, None, None]:
proto_topic_path = publisher_client.topic_path(PROJECT_ID, PROTO_TOPIC_ID)
try:
proto_topic = publisher_client.get_topic(
request={"topic": proto_topic_path})
except NotFound:
proto_topic = publisher_client.create_topic(
request={
"name": proto_topic_path,
"schema_settings": {
"schema": proto_schema,
"encoding": Encoding.BINARY,
},
})
yield proto_topic.name
publisher_client.delete_topic(request={"topic": proto_topic.name})
def publish(
self,
topic: str,
messages: List[Dict],
project_id: Optional[str] = None,
) -> None:
"""
Publishes messages to a Pub/Sub topic.
:param topic: the Pub/Sub topic to which to publish; do not
include the ``projects/{project}/topics/`` prefix.
:type topic: str
:param messages: messages to publish; if the data field in a
message is set, it should be a bytestring (utf-8 encoded)
:type messages: list of PubSub messages; see
http://cloud.google.com/pubsub/docs/reference/rest/v1/PubsubMessage
:param project_id: Optional, the GCP project ID in which to publish.
If set to None or missing, the default project_id from the GCP connection is used.
:type project_id: str
"""
if not project_id:
raise ValueError("Project ID should be set.")
self._validate_messages(messages)
publisher = self.get_conn()
topic_path = PublisherClient.topic_path(project_id, topic) # pylint: disable=no-member
self.log.info("Publish %d messages to topic (path) %s", len(messages),
topic_path)
try:
for message in messages:
publisher.publish(topic=topic_path,
data=message.get("data", b''),
**message.get('attributes', {}))
except GoogleAPICallError as e:
raise PubSubException(
'Error publishing to topic {}'.format(topic_path), e)
self.log.info("Published %d messages to topic (path) %s",
len(messages), topic_path)
def setUpClass(cls):
super().setUpClass()
cls.project = os.environ.get("PUBSUB_PROJECT", "testing")
cls.credentials = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS",
None)
for subscription in ["test0", "test1"]:
try:
publisher = PublisherClient()
publisher.create_topic(
publisher.topic_path(cls.project, subscription))
except AlreadyExists:
pass
try:
subscriber = SubscriberClient()
subscriber.create_subscription(
subscriber.subscription_path(cls.project, subscription),
subscriber.topic_path(cls.project, subscription))
except AlreadyExists:
pass
def avro_topic(publisher_client: pubsub_v1.PublisherClient,
avro_schema: str) -> Generator[str, None, None]:
from google.pubsub_v1.types import Encoding
avro_topic_path = publisher_client.topic_path(PROJECT_ID, AVRO_TOPIC_ID)
try:
avro_topic = publisher_client.get_topic(
request={"topic": avro_topic_path})
except NotFound:
avro_topic = publisher_client.create_topic(
request={
"name": avro_topic_path,
"schema_settings": {
"schema": avro_schema,
"encoding": Encoding.BINARY,
},
})
yield avro_topic.name
publisher_client.delete_topic(request={"topic": avro_topic.name})
def create_topic(
self,
topic: str,
project_id: str,
fail_if_exists: bool = False,
labels: Optional[Dict[str, str]] = None,
message_storage_policy: Union[Dict, MessageStoragePolicy] = None,
kms_key_name: Optional[str] = None,
retry: Optional[Retry] = None,
timeout: Optional[float] = None,
metadata: Optional[Sequence[Tuple[str, str]]] = None,
) -> None:
"""
Creates a Pub/Sub topic, if it does not already exist.
:param topic: the Pub/Sub topic name to create; do not
include the ``projects/{project}/topics/`` prefix.
:type topic: str
:param project_id: Optional, the Google Cloud project ID in which to create the topic
If set to None or missing, the default project_id from the Google Cloud connection is used.
:type project_id: str
:param fail_if_exists: if set, raise an exception if the topic
already exists
:type fail_if_exists: bool
:param labels: Client-assigned labels; see
https://cloud.google.com/pubsub/docs/labels
:type labels: Dict[str, str]
:param message_storage_policy: Policy constraining the set
of Google Cloud regions where messages published to
the topic may be stored. If not present, then no constraints
are in effect.
:type message_storage_policy:
Union[Dict, google.cloud.pubsub_v1.types.MessageStoragePolicy]
:param kms_key_name: The resource name of the Cloud KMS CryptoKey
to be used to protect access to messages published on this topic.
The expected format is
``projects/*/locations/*/keyRings/*/cryptoKeys/*``.
:type kms_key_name: str
:param retry: (Optional) A retry object used to retry requests.
If None is specified, requests will not be retried.
:type retry: google.api_core.retry.Retry
:param timeout: (Optional) The amount of time, in seconds, to wait for the request
to complete. Note that if retry is specified, the timeout applies to each
individual attempt.
:type timeout: float
:param metadata: (Optional) Additional metadata that is provided to the method.
:type metadata: Sequence[Tuple[str, str]]]
"""
publisher = self.get_conn()
topic_path = PublisherClient.topic_path(project_id, topic) # pylint: disable=no-member
# Add airflow-version label to the topic
labels = labels or {}
labels['airflow-version'] = 'v' + version.replace('.', '-').replace('+', '-')
self.log.info("Creating topic (path) %s", topic_path)
try:
# pylint: disable=no-member
publisher.create_topic(
name=topic_path,
labels=labels,
message_storage_policy=message_storage_policy,
kms_key_name=kms_key_name,
retry=retry,
timeout=timeout,
metadata=metadata,
)
except AlreadyExists:
self.log.warning('Topic already exists: %s', topic)
if fail_if_exists:
raise PubSubException('Topic already exists: {}'.format(topic))
except GoogleAPICallError as e:
raise PubSubException('Error creating topic {}'.format(topic), e)
self.log.info("Created topic (path) %s", topic_path)
