class PubSubConsumer:
GCLOUD_PUBSUB_PROJECT_ID = 'herbie-app'
def __init__(self):
self._subscriber = SubscriberClient()
def create_subscription(self, topic: str, subscription: str):
topic_path = self._subscriber.topic_path(self.GCLOUD_PUBSUB_PROJECT_ID,
topic)
subscription_path = self._subscriber.subscription_path(
self.GCLOUD_PUBSUB_PROJECT_ID, subscription)
try:
self._subscriber.create_subscription(subscription_path, topic_path)
except AlreadyExists as e:
print(f'Subscription {subscription} already exists.')
pass
def subscribe(self, subscription: str):
project_id = self.GCLOUD_PUBSUB_PROJECT_ID
subscription_path = self._subscriber.subscription_path(
project_id, subscription)
future = self._subscriber.subscribe(subscription_path,
self._subscribe_callback)
print(f'Listening for messages on {subscription_path}')
try:
future.result()
except KeyboardInterrupt:
future.cancel()
def _subscribe_callback(self, message):
print(message.data)
def setup_base_subscriber(subscriber: SubscriberClient,
subscriptions: Subscriptions,
**kwargs: Any) -> List[StreamingPullFuture]:
return [
subscriber.subscribe(
subscriber.subscription_path(GOOGLE_PROJECT, subscription_name),
callback=generate_callback(Message, message_handler),
**kwargs) for subscription_name, message_handler in subscriptions
]
class Subscriber(object):
def __init__(self, project_id, subscription_name, activity_id=None):
self.subscriber_client = SubscriberClient()
self.subscription_path = self.subscriber_client.subscription_path(
project_id, subscription_name)
self.activity_id = activity_id
self.user_report = user_report_handler.UserReportHandler()
self.db = sqlalchemy.create_engine(
sqlalchemy.engine.url.URL(drivername="mysql+pymysql",
username=_DB_USER,
password=_DB_PASS,
database=_DB_NAME))
# To run with local cloud proxy instance add query to the method above:
# query={"unix_socket": "/cloudsql/{}".format(_SQL_CONNECTION)}
def get_streaming_subscription(self):
"""Returns streaming subscription on Google Pub/Sub topic."""
subscription = self.subscriber_client.subscribe(
self.subscription_path, callback=self.get_callback())
return subscription
def _get_user_details_id(self, participant_id):
"""Returns the user_details_id for the given participant_id."""
with self.db.connect() as conn:
query = sqlalchemy.text(
'''SELECT user_details_id FROM participant_study_info
WHERE participant_id=:participant_id''')
result = conn.execute(query,
participant_id=participant_id).fetchone()
if not result:
logging.warning(
'Failed fetching user_id for participant {}'.format(
participant_id))
return None
return result[0]
@abc.abstractmethod
def get_callback(self):
"""Defines callback invoked by the subscription."""
class EmsSubscriberClient:
def __init__(self):
self.__client = SubscriberClient()
def subscribe(
self, subscription: str,
callback: Callable[[EmsMessage], None]) -> EmsStreamingFuture:
def callback_wrapper(message: Message) -> None:
callback(
EmsMessage(message.ack_id, message.data, message.attributes))
message.ack()
future = self.__client.subscribe(subscription=subscription,
callback=callback_wrapper)
return EmsStreamingFuture(future)
def pull(self,
subscription: str,
max_messages: int,
return_immediately: bool = None) -> Iterator[EmsMessage]:
messages = self.__client.api.pull(
subscription=subscription,
max_messages=max_messages,
return_immediately=return_immediately).received_messages
# pylint: disable=map-builtin-not-iterating
return map(
lambda msg: EmsMessage(msg.ack_id, msg.message.data, msg.message.
attributes), messages)
def acknowledge(self, subscription: str, ack_ids: List[str]) -> None:
self.__client.api.acknowledge(subscription=subscription,
ack_ids=ack_ids)
def create_subscription_if_not_exists(self, project_id: str,
topic_name: str,
subscription_name: str):
topic_path = self.__client.api.topic_path(project_id, topic_name)
subscription_path = self.__client.api.subscription_path(
project_id, subscription_name)
try:
self.__client.api.create_subscription(subscription_path,
topic_path)
LOGGER.info("Subscription %s created for topic %s in project %s",
subscription_name, topic_name, project_id)
except AlreadyExists:
LOGGER.info(
"Subscription %s already exists for topic %s in project %s",
subscription_name, topic_name, project_id)
def delete_subscription_if_exists(self, project_id: str,
subscription_name: str):
subscription_path = self.__client.api.subscription_path(
project_id, subscription_name)
try:
self.__client.api.delete_subscription(subscription_path)
LOGGER.info("Subscription %s deleted in project %s",
subscription_name, project_id)
except NotFound:
LOGGER.info("Subscription %s not found in project %s",
subscription_name, project_id)
