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:
- The python process does not end/exit until I interrupt it.
- 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.