Rename post to extract, implement Pub/Sub functionality in the ETL driver.

Author: Milo Hyben

etl/post/__init__.py renamed to etl/extract/__init__.py

@@ -19,7 +19,7 @@
 
 
 @functions_framework.http
-def etl_post(request: flask.Request):
+def etl_extract(request: flask.Request):
     """HTTP Cloud Function.
     Args:
         request (flask.Request): The request object.
@@ -59,12 +59,13 @@ def etl_post(request: flask.Request):
         'request_id': request_id,
         'timestamp': datetime.datetime.utcnow().isoformat(),
         'type': request.path,
-        'submitting_user': email_from_id_token(auth.token),
-        'body': jbody_str,
+        'submitting_user': email_from_id_token(auth.token)
     }
 
     # throw an exception if one occurs
-    errors = _BQ_CLIENT.insert_rows_json(BIGQUERY_TABLE, [bq_obj])
+    errors = _BQ_CLIENT.insert_rows_json(
+        BIGQUERY_TABLE, [bq_obj | {'body': jbody_str}]
+    )
     if errors:
         return {
             'success': False,
@@ -73,10 +74,13 @@ def etl_post(request: flask.Request):
         }, 500
 
     # publish to pubsub
-    # message contains all the attributes except body which can be large and already stored in BQ table
-    pb_obj = {k: v for k, v in bq_obj.items() if k not in ['body']}
+    # message contains all the attributes except body which can be large
+    # and already stored in BQ table
     try:
-        _PUBSUB_CLIENT.publish(PUBSUB_TOPIC, json.dumps(pb_obj).encode(), content_type='application/json')
+        _PUBSUB_CLIENT.publish(
+            PUBSUB_TOPIC,
+            json.dumps(bq_obj).encode()
+        )
     except Exception as e:  # pylint: disable=broad-exception-caught
         logging.error(f'Failed to publish to pubsub: {e}')
 

etl/post/main.py renamed to etl/extract/main.py

@@ -31,16 +31,24 @@ def etl_load(request: flask.Request):
         "request_id": "70eb6292-6311-44cf-9c9b-2b38bb076699"
     }
 
+    At the moment pulumi does not support unwrapping of push messages:
+    https://github.com/pulumi/pulumi-gcp/issues/1142
+
+    We need to support both
     """
 
     auth = request.authorization
     if not auth or not auth.token:
         return {'success': False, 'message': 'No auth token provided'}, 401
 
+    logging.info(f'auth {auth}')
+
     # if mimetype might not be set esp. when PubSub pushing from another topic,
     # try to force conversion and if fails just return None
     jbody = request.get_json(force=True, silent=True)
 
+    logging.info(f'jbody: {jbody}')
+
     if callable(jbody):
         # request.json is it in reality, but the type checker is saying it's callable
         jbody = jbody()

etl/post/requirements.txt renamed to etl/extract/requirements.txt

@@ -14,9 +14,11 @@
 from cpg_infra.utils import archive_folder
 
 # this gets moved around during the pip install
-ETL_FOLDER = Path(__file__).parent / 'etl'
+# ETL_FOLDER = Path(__file__).parent / 'etl'
+ETL_FOLDER = Path(__file__).parent.parent / 'etl'
 PATH_TO_ETL_BQ_SCHEMA = ETL_FOLDER / 'bq_schema.json'
-PATH_TO_ETL_ENDPOINT = ETL_FOLDER / 'endpoint'
+# PATH_TO_ETL_EXTRACT = ETL_FOLDER / 'extract'
+# PATH_TO_ETL_ENDPOINT = ETL_FOLDER / 'load'
 
 
 class MetamistInfrastructure(CpgInfrastructurePlugin):
@@ -103,7 +105,7 @@ def source_bucket(self):
         in a Google Cloud Storage bucket.
         """
         return gcp.storage.Bucket(
-            f'metamist-source-bucket',
+            'metamist-source-bucket',
             name=f'{self.config.gcp.dataset_storage_prefix}metamist-source-bucket',
             location=self.config.gcp.region,
             project=self.config.sample_metadata.gcp.project,
@@ -143,11 +145,97 @@ def etl_pubsub_topic(self):
         Pubsub topic to trigger the etl function
         """
         return gcp.pubsub.Topic(
-            f'metamist-etl-topic',
+            'metamist-etl-topic',
             project=self.config.sample_metadata.gcp.project,
             opts=pulumi.ResourceOptions(depends_on=[self._svc_pubsub]),
         )
 
+    @cached_property
+    def etl_pubsub_dead_letters_topic(self):
+        """
+        Pubsub dead_letters topic to capture failed jobs
+        """
+        topic = gcp.pubsub.Topic(
+            'metamist-etl-dead-letters-topic',
+            project=self.config.sample_metadata.gcp.project,
+            opts=pulumi.ResourceOptions(depends_on=[self._svc_pubsub]),
+        )
+
+        # give publisher permission to service account
+        gcp.pubsub.TopicIAMPolicy(
+            'metamist-etl-dead-letters-topic-iam-policy',
+            project=self.config.sample_metadata.gcp.project,
+            topic=topic.name,
+            policy_data=self.prepare_service_account_policy_data(
+                'roles/pubsub.publisher'
+            )
+        )
+
+        return topic
+
+    @cached_property
+    def etl_pubsub_push_subscription(self):
+        """
+        Pubsub push_subscription to topic, new messages to topic triggeres load process
+        """
+        subscription = gcp.pubsub.Subscription(
+            'metamist-etl-subscription',
+            topic=self.etl_pubsub_topic.name,
+            ack_deadline_seconds=20,
+            dead_letter_policy=gcp.pubsub.SubscriptionDeadLetterPolicyArgs(
+                dead_letter_topic=self.etl_pubsub_dead_letters_topic.id,
+                max_delivery_attempts=5,
+            ),
+            push_config=gcp.pubsub.SubscriptionPushConfigArgs(
+                push_endpoint=self.etl_load_function.service_config.uri,
+                oidc_token=gcp.pubsub.SubscriptionPushConfigOidcTokenArgs(
+                    service_account_email=self.etl_service_account.email,
+                ),
+                attributes={
+                    'x-goog-version': 'v1',
+                },
+            ),
+            project=self.config.sample_metadata.gcp.project,
+            opts=pulumi.ResourceOptions(
+                depends_on=[
+                    self._svc_pubsub,
+                    self.etl_pubsub_topic,
+                    self.etl_load_function,
+                    self.etl_pubsub_dead_letters_topic,
+                    self.etl_pubsub_dead_letter_subscription
+                ]
+            ),
+        )
+
+        # give subscriber permission to service account
+        gcp.pubsub.SubscriptionIAMPolicy(
+            'metamist-etl-pubsub-topic-subscription-policy',
+            project=self.config.sample_metadata.gcp.project,
+            subscription=subscription.name,
+            policy_data=self.prepare_service_account_policy_data(
+                'roles/pubsub.subscriber'
+            )
+        )
+
+        return subscription
+
+    @cached_property
+    def etl_pubsub_dead_letter_subscription(self):
+        """
+        Dead letter subscription
+        """
+        return gcp.pubsub.Subscription(
+            'metamist-etl-dead-letter-subscription',
+            topic=self.etl_pubsub_dead_letters_topic.name,
+            project=self.config.sample_metadata.gcp.project,
+            ack_deadline_seconds=20,
+            opts=pulumi.ResourceOptions(
+                depends_on=[
+                    self.etl_pubsub_dead_letters_topic
+                ]
+            ),
+        )
+
     @cached_property
     def etl_bigquery_dataset(self):
         """
@@ -219,7 +307,38 @@ def slack_channel(self):
             project=self.config.sample_metadata.gcp.project,
         )
 
+    def prepare_service_account_policy_data(self, role):
+        """
+        prepare_service_account_policy_data
+
+        Args:
+            role (_type_): _description_
+
+        Returns:
+            _type_: _description_
+        """
+        # get project
+        project = gcp.organizations.get_project()
+
+        return gcp.organizations.get_iam_policy(
+            bindings=[
+                gcp.organizations.GetIAMPolicyBindingArgs(
+                    role=role,
+                    members=[
+                        pulumi.Output.concat(
+                            'serviceAccount:service-',
+                            project.number,
+                            '@gcp-sa-pubsub.iam.gserviceaccount.com'
+                        )
+                    ],
+                )
+            ]
+        ).policy_data
+
     def setup_etl(self):
+        """
+        setup_etl
+        """
         # give the etl_service_account ability to write to bigquery
         gcp.bigquery.DatasetAccess(
             'metamist-etl-bq-dataset-access',
@@ -229,21 +348,80 @@ def setup_etl(self):
             user_by_email=self.etl_service_account.email,
         )
 
+        # give the etl_service_account ability to execute bigquery jobs
+        gcp.projects.IAMMember(
+            'metamist-etl-bq-job-user-role',
+            project=self.config.sample_metadata.gcp.project,
+            role='roles/bigquery.jobUser',
+            member=pulumi.Output.concat(
+                'serviceAccount:', self.etl_service_account.email
+            ),
+        )
+
+        # pubsub_v1.PublisherClient.publish User not authorized to perform this action
+        gcp.projects.IAMMember(
+            'metamist-etl-editor-role',
+            project=self.config.sample_metadata.gcp.project,
+            role='roles/editor',
+            member=pulumi.Output.concat(
+                'serviceAccount:', self.etl_service_account.email
+            ),
+        )
+
+        self.setup_etl_functions()
+        self.setup_etl_pubsub()
+
         self.setup_metamist_etl_accessors()
         self.setup_slack_notification()
 
+    def setup_etl_functions(self):
+        """
+        setup_etl_functions
+        """
+        self.etl_extract_function
+        self.etl_load_function
+
+    def setup_etl_pubsub(self):
+        """
+        setup_etl_pubsub
+        """
+        self.etl_pubsub_dead_letter_subscription
+        self.etl_pubsub_push_subscription
+
     @cached_property
-    def etl_function(self):
+    def etl_extract_function(self):
         """
-        Driver function to setup the etl infrastructure
+        etl_extract_function
+
+        Returns:
+            _type_: _description_
         """
+        return self.etl_function('extract')
+
+    @cached_property
+    def etl_load_function(self):
+        """
+        etl_load_function
+
+        Returns:
+            _type_: _description_
+        """
+        return self.etl_function('load')
+
+    def etl_function(self, f_name: str):
+        """
+        Driver function to setup the etl cloud function
+        """
+
+        path_to_func_folder = ETL_FOLDER / f_name
+
         # The Cloud Function source code itself needs to be zipped up into an
         # archive, which we create using the pulumi.AssetArchive primitive.
-        archive = archive_folder(str(PATH_TO_ETL_ENDPOINT.absolute()))
+        archive = archive_folder(str(path_to_func_folder.absolute()))
 
         # Create the single Cloud Storage object, which contains the source code
         source_archive_object = gcp.storage.BucketObject(
-            'metamist-etl-endpoint-source-code',
+            f'metamist-etl-{f_name}-source-code',
             # updating the source archive object does not trigger the cloud function
             # to actually updating the source because it's based on the name,
             # allow Pulumi to create a new name each time it gets updated
@@ -253,11 +431,11 @@ def etl_function(self):
         )
 
         fxn = gcp.cloudfunctionsv2.Function(
-            'metamist-etl-endpoint-source-code',
-            name='metamist-etl',
+            f'metamist-etl-{f_name}-source-code',
+            name=f'metamist-etl-{f_name}',
             build_config=gcp.cloudfunctionsv2.FunctionBuildConfigArgs(
                 runtime='python311',
-                entry_point='etl_post',
+                entry_point=f'etl_{f_name}',
                 environment_variables={},
                 source=gcp.cloudfunctionsv2.FunctionBuildConfigSourceArgs(
                     storage_source=gcp.cloudfunctionsv2.FunctionBuildConfigSourceStorageSourceArgs(
@@ -282,7 +460,8 @@ def etl_function(self):
                         self.etl_bigquery_table.table_id,
                     ),
                     'PUBSUB_TOPIC': self.etl_pubsub_topic.id,
-                    'ALLOWED_USERS': '[email protected]',
+                    # 'ALLOWED_USERS': '[email protected]',
+                    'ALLOWED_USERS': '[email protected]',
                 },
                 ingress_settings='ALLOW_ALL',
                 all_traffic_on_latest_revision=True,