Merge pull request #59 from google/fix_local_runner

Fix DirectRunner

Author: diogoaihara

megalista_dataflow/main.py

@@ -16,6 +16,7 @@
 import warnings
 
 import apache_beam as beam
+from apache_beam import coders
 from apache_beam.options.pipeline_options import PipelineOptions
 from mappers.ads_user_list_pii_hashing_mapper import \
     AdsUserListPIIHashingMapper
@@ -24,7 +25,7 @@
 from models.oauth_credentials import OAuthCredentials
 from models.options import DataflowOptions
 from models.sheets_config import SheetsConfig
-from sources.batches_from_executions import BatchesFromExecutions
+from sources.batches_from_executions import BatchesFromExecutions, ExecutionCoder
 from sources.primary_execution_source import PrimaryExecutionSource
 from uploaders.big_query.transactional_events_results_writer import TransactionalEventsResultsWriter
 from uploaders.campaign_manager.campaign_manager_conversion_uploader import CampaignManagerConversionUploaderDoFn
@@ -272,6 +273,8 @@ def run(argv=None):
 
     params = MegalistaStepParams(oauth_credentials, dataflow_options)
 
+    coders.registry.register_coder(Execution, ExecutionCoder)
+
     with beam.Pipeline(options=pipeline_options) as pipeline:
         executions = pipeline | "Load executions" >> beam.io.Read(execution_source)
 

megalista_dataflow/models/execution.py

@@ -80,6 +80,25 @@ def campaign_manager_account_id(self) -> str:
     def app_id(self) -> str:
         return self._app_id
 
+    def to_dict(self):
+        return {
+            'google_ads_account_id' : self.google_ads_account_id,
+            'mcc': self.mcc,
+            'google_analytics_account_id': self.google_analytics_account_id,
+            'campaign_manager_account_id': self.campaign_manager_account_id,
+            'app_id': self.app_id,
+        }
+
+    @staticmethod
+    def from_dict(dict_account_config):
+        return AccountConfig(
+            dict_account_config['google_ads_account_id'],
+            dict_account_config['mcc'],
+            dict_account_config['google_analytics_account_id'],
+            dict_account_config['campaign_manager_account_id'],
+            dict_account_config['app_id'],
+        )
+
     def __str__(self) -> str:
         return (
             f"\n[Account Config]\n\t"
@@ -129,6 +148,21 @@ def source_type(self) -> SourceType:
     def source_metadata(self) -> List[str]:
         return self._source_metadata
 
+    def to_dict(self):
+        return {
+            'source_name': self.source_name,
+            'source_type' : self.source_type.name,
+            'source_metadata': self.source_metadata,
+        }
+
+    @staticmethod
+    def from_dict(dict_source):
+        return Source(
+            dict_source['source_name'],
+            SourceType[dict_source['source_type']],
+            dict_source['source_metadata']
+        )
+
     def __eq__(self, other):
         return (
             self.source_name == other.source_name
@@ -170,6 +204,21 @@ def destination_type(self) -> DestinationType:
     def destination_metadata(self) -> List[str]:
         return self._destination_metadata
 
+    def to_dict(self):
+        return {
+            'destination_name': self.destination_name,
+            'destination_type': self.destination_type.name,
+            'destination_metadata': self.destination_metadata,
+        }
+
+    @staticmethod
+    def from_dict(dict_destination):
+        return Destination(
+            dict_destination['destination_name'],
+            DestinationType[dict_destination['destination_type']],
+            dict_destination['destination_metadata'],
+        )
+
     def __eq__(self, other) -> bool:
         return bool(
             self.destination_name == other.destination_name
@@ -206,6 +255,21 @@ def destination(self) -> Destination:
     def account_config(self) -> AccountConfig:
         return self._account_config
 
+    def to_dict(self):
+        return {
+            'account_config': self.account_config.to_dict(),
+            'source': self.source.to_dict(),
+            'destination': self.destination.to_dict()
+        }
+
+    @staticmethod
+    def from_dict(dict_execution):
+        return Execution(
+            AccountConfig.from_dict(dict_execution['account_config']),
+            Source.from_dict(dict_execution['source']),
+            Destination.from_dict(dict_execution['destination']),
+        )
+
     def __str__(self):
         return f"Origin name: {self.source.source_name}. Action: {self.destination.destination_type}. Destination name: {self.destination.destination_name}"
 

megalista_dataflow/requirements.txt

@@ -4,8 +4,8 @@ protobuf==3.13.0
 google-api-python-client==1.12.8
 google-cloud-core==1.4.1
 google-cloud-bigquery==1.27.2
-apache-beam[gcp]==2.34.0
-apache-beam==2.34.0
+apache-beam[gcp]==2.36.0
+apache-beam==2.36.0
 google-cloud-datastore==1.13.1
 google-apitools==0.5.31
 pytest==5.4.3

megalista_dataflow/sources/batches_from_executions.py

@@ -12,16 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, List, Iterable
-
 import apache_beam as beam
 import logging
+import json
 
+from apache_beam.coders import coders
 from apache_beam.options.value_provider import ValueProvider
 from google.cloud import bigquery
-from apache_beam.io.gcp.bigquery import ReadFromBigQueryRequest
-
 from models.execution import DestinationType, Execution, Batch
+from typing import Any, List, Iterable, Tuple, Dict
+
 
 _BIGQUERY_PAGE_SIZE = 20000
 
@@ -35,29 +35,43 @@ def _convert_row_to_dict(row):
     return dict
 
 
+class ExecutionCoder(coders.Coder):
+    """A custom coder for the Execution class."""
+
+    def encode(self, o):
+        return json.dumps(o.to_dict()).encode('utf-8')
+
+    def decode(self, s):
+        return Execution.from_dict(json.loads(s.decode('utf-8')))
+
+    def is_deterministic(self):
+        return True
+
+
 class BatchesFromExecutions(beam.PTransform):
     """
     Filter the received executions by the received action,
     load the data using the received source and group by that batch size and Execution.
     """
 
     class _ExecutionIntoBigQueryRequest(beam.DoFn):
-        def process(self, execution: Execution) -> Iterable[ReadFromBigQueryRequest]:
+
+        def process(self, execution: Execution) -> Iterable[Tuple[Execution, Dict[str, Any]]]:
             client = bigquery.Client()
             table_name = execution.source.source_metadata[0] + '.' + execution.source.source_metadata[1]
             table_name = table_name.replace('`', '')
             query = f"SELECT data.* FROM `{table_name}` AS data"
             logging.getLogger(_LOGGER_NAME).info(f'Reading from table {table_name} for Execution {execution}')
             rows_iterator = client.query(query).result(page_size=_BIGQUERY_PAGE_SIZE)
             for row in rows_iterator:
-                yield {'execution': execution, 'row': _convert_row_to_dict(row)}
+                yield execution, _convert_row_to_dict(row)
 
     class _ExecutionIntoBigQueryRequestTransactional(beam.DoFn):
 
         def __init__(self, bq_ops_dataset):
             self._bq_ops_dataset = bq_ops_dataset
 
-        def process(self, execution: Execution) -> Iterable[ReadFromBigQueryRequest]:
+        def process(self, execution: Execution) -> Iterable[Tuple[Execution, Dict[str, Any]]]:
             table_name = execution.source.source_metadata[0] + \
                 '.' + execution.source.source_metadata[1]
             table_name = table_name.replace('`', '')
@@ -86,7 +100,7 @@ def process(self, execution: Execution) -> Iterable[ReadFromBigQueryRequest]:
                 f'Reading from table {table_name} for Execution {execution}')
             rows_iterator = client.query(query).result(page_size=_BIGQUERY_PAGE_SIZE)
             for row in rows_iterator:
-                yield {'execution': execution, 'row': _convert_row_to_dict(row)}
+                yield execution, _convert_row_to_dict(row)
 
 
     class _BatchElements(beam.DoFn):
@@ -102,7 +116,7 @@ def process(self, grouped_elements):
                 if i != 0 and i % self._batch_size == 0:
                     yield Batch(execution, batch)
                     batch = []
-                batch.append(element['row'])
+                batch.append(element)
             yield Batch(execution, batch)
 
     def __init__(
@@ -131,6 +145,6 @@ def expand(self, executions):
             executions
             | beam.Filter(lambda execution: execution.destination.destination_type == self._destination_type)
             | beam.ParDo(self._get_bq_request_class())
-            | beam.GroupBy(lambda x: x['execution'])
+            | beam.GroupByKey()
             | beam.ParDo(self._BatchElements(self._batch_size))
         )