Fail to stream data from Kafka to BigQuery with Apache Beam

Viewed 70

I am trying to load a stream of data from a kafka topic to a BigQuery table. While I am able to source the stream from the kafka topic and do transformations on it, I am stuck on loading the transformed data to a BQ table.
Note: I am using apache beam's python SDK with direct runner (right now) to test things. Here's the code:

import os
import argparse
import json
import logging

import apache_beam as beam
from apache_beam.io.gcp.bigquery import WriteToBigQuery
from beam_nuggets.io import kafkaio


def run(argv=None):

    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--bq_table',
        required=True,
        help=('Output BigQuery table for results specified as: '
              'PROJECT:DATASET.TABLE'))

    known_args, pipeline_args = parser.parse_known_args(argv)

    bootstrap_server = "localhost:9092"
    kafka_topic = 'some_topic'
    consumer_config = {
        'bootstrap_servers': bootstrap_server,
        'group_id': 'etl',
        'topic': kafka_topic,
        'auto_offset_reset': 'earliest'
    }

    with beam.Pipeline(argv=pipeline_args) as p:
        events = (
            p | 'Read from kafka' >> kafkaio.KafkaConsume(
                consumer_config=consumer_config, value_decoder=bytes.decode)
            | 'toDict' >> beam.MapTuple(lambda k, v: json.loads(v))
            | 'extract' >> beam.Map(lambda e: {'x': e['key1'], 'y': e['key2']})
            # | 'log' >> beam.ParDo(logging.info)   # if I uncomment this (for validation), I see data printed in my console log
        )
        _ = events | 'Load data to BQ' >> WriteToBigQuery(known_args.bq_table,
                                                          schema='x:STRING, y:STRING',
                                                          write_disposition=beam.io.BigQueryDisposition.WRITE_APPEND,
                                                          create_disposition=beam.io.BigQueryDisposition.CREATE_IF_NEEDED,
                                                          ignore_unknown_columns=True, method='STREAMING_INSERTS')

        p.run().wait_until_finish()


if __name__ == "__main__":

    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "/path/to/credentials/file.json"
    logging.getLogger().setLevel(logging.DEBUG)
    logging.getLogger("kafka").setLevel(logging.INFO)
    run()

And this is how I run the above code:

python3 main.py \
--bq_table=<table_id> \
--streaming \
--runner=DirectRunner

I have tried using the batch mode to data insertion as well with WriteToBigQuery (method=FILE_LOADS) and providing a temp GCS location, but that didn't help either.
There is no error even though I have enabled debug logs. So, it's getting really difficult to trace the issue to its source. What do you think I am missing?

Edit:

  1. The python process does not end/exit until I interrupt it.
  2. The Kafka consumer group lag is 0, which indicates that the data is being fetched and processed but not getting loaded to the BQ table.
0 Answers
Related