#!/usr/bin/env python
"""Extract events from kafka and write them to hdfs
"""
import json
from pyspark.sql import SparkSession, Row
from pyspark.sql.functions import udf


@udf('boolean')
def is_purchase(event_as_json):
    event = json.loads(event_as_json)
    if event['event_type'] == 'purchase_sword':
        return True
    return False


def main():
    """main
    """
    spark = SparkSession \
        .builder \
        .appName("ExtractEventsJob") \
        .getOrCreate()

    raw_events = spark \
        .read \
        .format("kafka") \
        .option("kafka.bootstrap.servers", "kafka:29092") \
        .option("subscribe", "events") \
        .option("startingOffsets", "earliest") \
        .option("endingOffsets", "latest") \
        .load()

    purchase_events = raw_events \
        .select(raw_events.value.cast('string').alias('raw'),
                raw_events.timestamp.cast('string')) \
        .filter(is_purchase('raw'))

    extracted_purchase_events = purchase_events \
        .rdd \
        .map(lambda r: Row(timestamp=r.timestamp, **json.loads(r.raw))) \
        .toDF()
    extracted_purchase_events.printSchema()
    extracted_purchase_events.show()


if __name__ == "__main__":
    main()

[((9, 1, 9, 15), 'pyspark.sql.functions.udf', 'udf', ({(9, 5, 9, 14): '"""boolean"""'}, {}), "('boolean')", False, 'from pyspark.sql.functions import udf\n'), ((11, 12, 11, 37), 'json.loads', 'json.loads', ({(11, 23, 11, 36): 'event_as_json'}, {}), '(event_as_json)', False, 'import json\n'), ((20, 12, 22, 36), 'pyspark.sql.SparkSession.builder.appName', 'SparkSession.builder.appName', ({(22, 17, 22, 35): '"""ExtractEventsJob"""'}, {}), "('ExtractEventsJob')", False, 'from pyspark.sql import SparkSession, Row\n'), ((41, 52, 41, 69), 'json.loads', 'json.loads', ({(41, 63, 41, 68): 'r.raw'}, {}), '(r.raw)', False, 'import json\n')]