main.py

# Copyright 2020 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""This script demonstrates the concept of referencing multiple forms of side input data in GCS with main input data
loaded from pubsub, and triggering the refresh of side input data via PubSub notification on demand

Description of main and side inputs with sample records:
-------------------------------------------------------
Main Input: Sales event with schema:
  {"Txid":1, "productname": "Product 5 XL", "qty": 1, "sales": 97.65, "event_timestamp": "2020-05-02 19:08:24.796892"}

Example contains 3 side inputs:  Bonus points, Discount percent and Category
1. Bonus points -> {"productname": "Product 2", "bonuspoints": 204}
2. Discount percent -> {"productname": "Product 2", "discountpct": 0.62}
3. Category -> {"productname": "Product 3", "category": "laptop"}

Args:
    --project: GCP project ID
    --sideinput-notification-subscription: pub sub subscription name notifying new base path from which the side input
        data needs to be reloaded.
    --maininput-subscription: pub sub subscription name from which main input (i.e sales events) are consumed
    --sideinput-window-duration:  Side Input window duration. Ideally Global window needs to be used instead of fixed
      window. However on refresh new state is not replacing old state. Hence Fixed windows of longer duration (180 days
      to an year) can be used to avoid expiration of window. Duration must be provided in seconds.
    --maininput-window-duration: Main Input window duration in seconds
    (--demorun | --outputpath=<OUTPUTPATH>): Use demo run if ouptput should be written to log/stdout for inspection and
            specify output path for writing to persistent storage .


Example Usage:
python -m sideinput_refresh --project=<project> --region=us-central1  --runner=DataFlowRunner \
 --sideinput-notification-subscription=<subscription-name> \
 --maininput-subscription=<subscription-name> \
 --sideinput-window-duration=<518400> \
 --maininput-window-duration=<300> \
 (--demorun | --outputpath=<GCS path>) \
 --staging_location=<GCS path> \
 --temp_location=<GCS Path>
"""

import argparse
import json
import logging
import re
from typing import Dict, List, Tuple

import apache_beam as beam
from apache_beam.options.pipeline_options import (GoogleCloudOptions,
                                                  PipelineOptions, SetupOptions,
                                                  StandardOptions)
from apache_beam.pvalue import AsDict

from sideinput_refresh import dofns, transforms


def get_subscription_path(project: str, subscription_name: str) -> str:
    """"Helper function to build topic subscription path from project and subscription name

      Args:
        project: GCP Project id
        subscription_name: PubSub Subscription name

      Returns:
        PubSub Subscription Path
  """
    if re.match('projects/([^/]+)/subscriptions/(.+)', subscription_name):
        return subscription_name

    return f"projects/{project}/subscriptions/{subscription_name}"


def parse_cliargs(argv: List[str]) -> Tuple[argparse.Namespace, List[str]]:
    """"Parses Command line arguments

      Args:
          argv: List of command line arguments

      Returns:
          Tuple containing Namespace object with specified arguments and list of extra arguments
  """
    args_parser = argparse.ArgumentParser()

    args_parser.add_argument('--project',
                             type=str,
                             help="Google Cloud project id",
                             required=True)

    args_parser.add_argument(
        '--sideinput-notification-subscription',
        type=str,
        help="Cloud Pub/Sub Side input notificaton subscription name",
        required=True)

    args_parser.add_argument('--maininput-subscription',
                             type=str,
                             help="Cloud Pub/Sub main input subscription name",
                             required=True)

    args_parser.add_argument(
        '--sideinput-window-duration',
        type=int,
        help="Duration of the side input window in seconds",
        required=True)

    args_parser.add_argument(
        '--maininput-window-duration',
        type=int,
        help="Duration of the main input window in seconds",
        required=True)

    sink_choices = args_parser.add_mutually_exclusive_group(required=True)

    sink_choices.add_argument(
        '--demorun',
        help="Flag to write the output to log sink instead of file sink",
        action='store_true')

    sink_choices.add_argument(
        '--outputpath',
        type=str,
        help="Output file path e.g: (gs://<buckername>/prefix)",
        action='store')

    return args_parser.parse_known_args(argv)


def get_sideinput_collections(sideinput_filepath: beam.pvalue.PCollection, readTransform: beam.PTransform) \
        -> Dict[str,beam.pvalue.PCollection]:
    """"Load Side Input data from respective file paths

      Args:
          sideinput_filepath: File path representing base path for side inputs to be loaded
          readTransform: Transform responsible for loading the side input data

      Returns:
          Dictionary containing Side Input name as key and corresponding PCollection as value
  """
    sideinput_types = get_sideinput_types()
    # yapf: disable
    filepaths =  (sideinput_filepath
                  | "Get side input paths from base path" >> beam.ParDo(
                     dofns.SplitToMultiple(sideinput_types)).with_outputs(*sideinput_types)
                  )

    sideinput_collections = {}
    for sideinput_type in sideinput_types:
        sideinput_collections[sideinput_type] = (filepaths[sideinput_type]
                                                    | f"Read {sideinput_type}" >> readTransform
                                                   | f"{sideinput_type}:Extract KV" >> beam.Map(
                                                     transforms.kv_of, "productname", sideinput_type)
                                                 )
    # yapf: enable
    return sideinput_collections



def get_enriched_events(salesevent: beam.pvalue.PCollection,sideinput_collections: Dict[str,beam.pvalue.PCollection]) \
        -> beam.pvalue.PCollection:
    """Gets enriched events by
        a) Call a transform that combining primary event with corresponding side input values
        b) Group events by dummy key to combine all events in a window into one shard
        c) Discard dummy key

     Args:
        salesevent: Event representing sales transaction
        sideinput_collections: Set of Side Input Collections
    """
    # yapf: disable
    return (salesevent
             | "Enrich event" >> beam.Map(transforms.enrich_event,
                                       AsDict(sideinput_collections["bonuspoints"]),
                                       AsDict(sideinput_collections["discountpct"]),
                                       AsDict(sideinput_collections["category"]))
             | "Group events by dummy Key" >> beam.GroupByKey()
             | "Discard dummy Key" >> beam.Values()
          )
    # yapf: enable


def get_sideinput_types() -> List[str]:
    """Returns list of side input types

     Returns:
        Side Input types
     """

    return ["bonuspoints", "discountpct", "category"]


def run(argv: List[str] = None, save_main_session: bool = True) -> None:
    """"Driver function to parse command line args, build pipeline options and invokes pipeline

      Args:
          argv: Command line arguments specified by the user
          save_main_session: Setting to indicate the visibility of modules to beam

      Returns:
          None
  """
    known_args, extra_args = parse_cliargs(argv)
    known_args.maininput_subscription = get_subscription_path(
        known_args.project, known_args.maininput_subscription)
    known_args.sideinput_notification_subscription = get_subscription_path(
        known_args.project, known_args.sideinput_notification_subscription)

    pipeline_options = PipelineOptions(flags=extra_args)
    pipeline_options.view_as(GoogleCloudOptions).project = known_args.project
    pipeline_options.view_as(SetupOptions).save_main_session = save_main_session
    pipeline_options.view_as(StandardOptions).streaming = True

    pipeline = beam.Pipeline(options=pipeline_options)

    # Get Side Input Refresh Notification
    # yapf: disable
    sideinput_filepath = (pipeline
                     | "Read Side input file path from pubsub subscription" >> beam.io.ReadFromPubSub(subscription=known_args.sideinput_notification_subscription)
                     | "Side input fixed window with early trigger" >> beam.WindowInto(
                    beam.window.FixedWindows(known_args.sideinput_window_duration),
                    trigger=beam.trigger.Repeatedly(beam.trigger.AfterCount(1)),
                    accumulation_mode=beam.trigger.AccumulationMode.DISCARDING)
               )

    # Get Side Input data
    # Ideally instead of writing transform to ReadSideInputData, an inbuilt transform textio.ReadAllFromText(coder=coders.UTF8JsonCoder())
    # can be used. However when ReadAllFromText is used despite of loading all data certain elements are not propogating well as side input
    # Created a BEAM issue https://issues.apache.org/jira/browse/BEAM-10148 for further investigation
    # Once issue is fixed transforms.ReadSideInputData() will be replaced with apache_beam.io.textio.ReadAllFromText(coder=coders.UTF8JsonCoder()))
    sideinput_collections = get_sideinput_collections(sideinput_filepath,transforms.ReadSideInputData())

    # Get Sales as main input
    primary_events = (pipeline
                       | "Read sales events from pubsub subscription" >> beam.io.ReadFromPubSub(subscription=known_args.maininput_subscription)
                       | "Main input fixed window" >> beam.WindowInto(beam.window.FixedWindows(known_args.maininput_window_duration))
                    )

    # Enrich primary events with side input data
    enriched_events = get_enriched_events(primary_events, sideinput_collections)

    # Write enriched events to sink
    if known_args.demorun:
        enriched_events | "Write output to Log" >> transforms.LogEvents()
    else:
        (enriched_events
                |  "Unpack the events" >> beam.FlatMap(lambda x: list(map(json.dumps, x)))
                |  "Write output to GCS" >> beam.io.fileio.WriteToFiles(known_args.outputpath, shards=1)
         )
    # yapf: enable

    pipeline_result = pipeline.run()

    # Used while testing locally
    if pipeline_options.view_as(StandardOptions).runner == "DirectRunner":
        pipeline_result.wait_until_finish()


if __name__ == "__main__":
    logging.getLogger().setLevel(level=logging.INFO)
    run()