[AIRFLOW-5477] Rewrite Google PubSub Hook to Google Cloud Python #6096

TobKed · 2019-09-13T13:45:57Z

Make sure you have checked all steps below.

Jira

My PR addresses the following Airflow Jira issues and references them in the PR title. For example, "[AIRFLOW-XXX] My Airflow PR"
- https://issues.apache.org/jira/browse/AIRFLOW-5477
- In case you are fixing a typo in the documentation you can prepend your commit with [AIRFLOW-XXX], code changes always need a Jira issue.
- In case you are proposing a fundamental code change, you need to create an Airflow Improvement Proposal (AIP).
- In case you are adding a dependency, check if the license complies with the ASF 3rd Party License Policy.

Description

Here are some details about my PR, including screenshots of any UI changes:

Tests

My PR adds the following unit tests OR does not need testing for this extremely good reason:

Commits

My commits all reference Jira issues in their subject lines, and I have squashed multiple commits if they address the same issue. In addition, my commits follow the guidelines from "How to write a good git commit message":
1. Subject is separated from body by a blank line
2. Subject is limited to 50 characters (not including Jira issue reference)
3. Subject does not end with a period
4. Subject uses the imperative mood ("add", not "adding")
5. Body wraps at 72 characters
6. Body explains "what" and "why", not "how"

Documentation

In case of new functionality, my PR adds documentation that describes how to use it.
- All the public functions and the classes in the PR contain docstrings that explain what it does
- If you implement backwards incompatible changes, please leave a note in the Updating.md so we can assign it to a appropriate release

codecov-io · 2019-09-13T22:36:41Z

Codecov Report

Merging #6096 into master will decrease coverage by 0.55%.
The diff coverage is 92.59%.

@@            Coverage Diff             @@
##           master    #6096      +/-   ##
==========================================
- Coverage   80.01%   79.45%   -0.56%     
==========================================
  Files         610      610              
  Lines       35176    35286     +110     
==========================================
- Hits        28147    28038     -109     
- Misses       7029     7248     +219

Impacted Files	Coverage Δ
airflow/gcp/example_dags/example_pubsub.py	`0% <0%> (ø)`	⬆️
airflow/gcp/sensors/pubsub.py	`94.28% <83.33%> (-5.72%)`	⬇️
airflow/gcp/operators/pubsub.py	`90.69% <83.78%> (-9.31%)`	⬇️
airflow/gcp/hooks/pubsub.py	`99.33% <99.22%> (+5.58%)`	⬆️
airflow/operators/postgres_operator.py	`0% <0%> (-100%)`	⬇️
airflow/operators/mysql_operator.py	`0% <0%> (-100%)`	⬇️
airflow/operators/mysql_to_hive.py	`0% <0%> (-100%)`	⬇️
airflow/operators/generic_transfer.py	`0% <0%> (-100%)`	⬇️
airflow/executors/celery_executor.py	`40.74% <0%> (-35.56%)`	⬇️
airflow/utils/log/wasb_task_handler.py	`32.87% <0%> (-9.59%)`	⬇️
... and 14 more

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 781d001...9db5c94. Read the comment docs.

airflow/gcp/hooks/pubsub.py

mik-laj · 2019-09-14T11:45:32Z

airflow/gcp/hooks/pubsub.py

+            )
+        except NotFound:
+            message = 'Topic does not exist: {}'.format(topic_path)
+            self.log.warning(message)


We should avoid formatting string before passing to logger.
Reference
https://github.com/apache/airflow/pull/4804/files
#5681

airflow/gcp/hooks/pubsub.py

mik-laj · 2019-09-14T11:53:34Z

airflow/gcp/hooks/pubsub.py

+                    data=message.get("data", b''),
+                    **message.get('attributes', {})
+                )
+        except GoogleAPICallError as e:
            raise PubSubException(


This part of code is not covered by unit test. Can you complete missing unit tests?

I've added test.

airflow/gcp/hooks/pubsub.py

mik-laj · 2019-09-14T11:55:29Z

airflow/gcp/hooks/pubsub.py

+            self.log.warning(message)
+            if fail_if_exists:
+                raise PubSubException(message)
+        except GoogleAPICallError as e:


This part of code is not covered by unit test. Can you complete missing unit tests?

I've added test.

airflow/gcp/hooks/pubsub.py

mik-laj · 2019-09-14T11:56:01Z

airflow/gcp/hooks/pubsub.py

        except HttpError as e:
            raise PubSubException(
                'Error pulling messages from subscription {}'.format(
-                    full_subscription), e)
+                    subscription_path), e)
+        except GoogleAPICallError as e:


This part of code is not covered by unit test. Can you complete missing unit tests?

I've added test.

airflow/gcp/hooks/pubsub.py

TobKed · 2019-09-16T14:11:33Z

@mik-laj I've applied changes. Please verify is correct.

mik-laj · 2019-09-17T10:54:38Z

@TobKed Today we have had various conversations about this integration. Is this PR done?

TobKed · 2019-09-17T12:12:38Z

@mik-laj Yes, hopefully PR is done. I've applied all necessary patches.

mik-laj · 2019-09-17T12:27:27Z

airflow/gcp/operators/pubsub.py

    """
-    template_fields = ['project', 'subscription']
+    template_fields = ['project_id', 'subscription']


Suggested change

template_fields = ['project_id', 'subscription']

template_fields = ['project', 'project_id', 'subscription']

I had to revert these changes (only "project" shall remain) because of the errors due to errors:
File "/workspace/airflow/models/baseoperator.py", line 662, in _do_render_template_fields content = getattr(parent, attr_name) AttributeError: 'PubSubPublishOperator' object has no attribute 'project'

Should be OK. Maybe (nice to have) we could add a warning in deprecation message, that templates won't work ?

mik-laj · 2019-09-17T12:28:55Z

airflow/gcp/operators/pubsub.py

    """
-    template_fields = ['project', 'topic', 'messages']
+    template_fields = ['project_id', 'topic', 'messages']


Suggested change

template_fields = ['project_id', 'topic', 'messages']

template_fields = ['project', 'project_id', 'topic', 'messages']

I had to revert these changes (only "project" shall remain) because of the errors due to errors:
File "/workspace/airflow/models/baseoperator.py", line 662, in _do_render_template_fields content = getattr(parent, attr_name) AttributeError: 'PubSubPublishOperator' object has no attribute 'project'

mik-laj · 2019-09-18T23:39:58Z

Can you do rebase?

TobKed · 2019-09-19T06:17:56Z

Rebased on the last master.

mik-laj · 2019-09-20T09:06:22Z

Everything looks good, but can you do rebase?

TobKed · 2019-09-20T14:10:49Z

I've rebased but due to fact than previous commits impacted PubSub significantly I had to make some changes (remove project from template fields, update PubSubPullSensor and adjust unit and system tests). Please take a look on last fixup.

TobKed · 2019-09-20T19:59:35Z

I've updated TestPubSubPullSensor (hook mocking).

turbaszek · 2019-09-21T10:43:37Z

airflow/gcp/sensors/pubsub.py

+        )
+
+        self._messages = [MessageToDict(m) for m in pulled_messages]
+
        if self._messages and self.ack_messages:
            if self.ack_messages:


Suggested change

if self.ack_messages:

This is a double assertion.

Good point! Thanks, I've fixed it.

TobKed · 2019-09-23T08:05:59Z

Updated due to comments, slightly improved test, rebased on last master.

TobKed · 2019-09-26T13:11:36Z

cc @mik-laj @potiuk

potiuk

Just one thing - It seems that the messages should not be base64 encoded when using the client (see example here: https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/pubsub/cloud-client/publisher.py). This means that in all tests, it should be fixed and comments should be updated. Also UPDATING.Md should mention that this should be changed in the input data (or maybe even we should somehow provide backwards compatibility and detect if someone passes a string rather than bytes and decode/encode appropriately).

potiuk · 2019-09-26T13:33:51Z

airflow/gcp/hooks/pubsub.py

+        publisher = self.get_conn()
+        topic_path = PublisherClient.topic_path(project_id, topic)  # pylint: disable=no-member
+
+        # TODO validation of messages


Should we leave that TODO here? I think we might already check it here simply to check if we can 'base64' decode What kind of validation we think we miss here? I think we already handle case where "data" is missing. I believe the publish method must check it already if the data is correctly encoded. If so, we should rather remove this TODO. Otherwise we should check ourselves what is the type of "data" field in the loop. I believe

isinstance(data, bytes)

for python 3+ should work fine.

I think this comment should not be present here.
I have separate PR for message validation waiting in queue:
https://github.com/PolideaInternal/airflow/pull/304/commits

In the last fixup I've included validation and deprecation warning.

potiuk · 2019-09-26T13:34:49Z

airflow/gcp/hooks/pubsub.py

        """
        Publishes messages to a Pub/Sub topic.

-        :param project: the GCP project ID in which to publish
-        :type project: str
        :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 already be base64 encoded.


I think this comment is wrong now. According to https://cloud.google.com/pubsub/docs/publisher (and the code below) the data parameter should be bytestring (utf-8 encoded) rather than base64 encoded string

You are right. Great point! I've updated docstrings.

potiuk · 2019-09-26T13:37:51Z

airflow/gcp/hooks/pubsub.py

        :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 GCP project ID in which to create the topic


Cool with all those params below. Looks like fully-featured implementation now 👍

potiuk · 2019-09-26T13:38:20Z

airflow/gcp/hooks/pubsub.py

+
+        # Add airflow-version label to the topic
+        labels = labels or {}
+        labels['airflow-version'] = 'v' + version.replace('.', '-').replace('+', '-')


potiuk · 2019-09-26T13:39:07Z

airflow/gcp/hooks/pubsub.py

        ack_deadline_secs: int = 10,
        fail_if_exists: bool = False,
+        push_config: Optional[Union[Dict, PushConfig]] = None,


Same here 👍

potiuk · 2019-09-26T13:40:51Z

airflow/gcp/operators/pubsub.py

            *args,
            **kwargs) -> None:
+


Nice that we keep backwards compatibility.

potiuk · 2019-09-26T13:43:12Z

airflow/gcp/operators/pubsub.py

    """
-    template_fields = ['project', 'subscription']
+    template_fields = ['project_id', 'subscription']


Should be OK. Maybe (nice to have) we could add a warning in deprecation message, that templates won't work ?

…ud Python

Co-Authored-By: Kamil Breguła <[email protected]>

TobKed · 2019-09-30T12:24:45Z

@potiuk I've added message validation. Now when user pass message data as a string which can be base64 decoded, an additional deprecation warning will be displayed, afterwardsPubSub exception will be raised. I've updated tests, added tests for message validation and added additional info in UPDATING.md.
About deprecation warning for template_fields I think it is possible however in my opinion it will introduce unnecessary amount of complication. I've included Information about template_filelds in UPDATING.md as well.

mik-laj added the provider:google Google (including GCP) related issues label Sep 13, 2019

TobKed force-pushed the pubsub-hook-with-google-api-python branch from 123eb84 to e019de8 Compare September 13, 2019 17:55

mik-laj force-pushed the pubsub-hook-with-google-api-python branch from e019de8 to 61a8117 Compare September 13, 2019 19:18