-
Notifications
You must be signed in to change notification settings - Fork 0
/
avro-data-producer.py
96 lines (80 loc) · 3.12 KB
/
avro-data-producer.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
from googlefinance import getQuotes
from kafka.errors import KafkaError, KafkaTimeoutError
import argparse
import atexit
import datetime
import logging
import json
import random
import schedule
import time
import confluent_kafka
from confluent_kafka import avro
from confluent_kafka.avro import AvroProducer
# - default kafka topic to write to
topic_name = 'stock-analyzer'
# - default kafka broker location
kafka_broker = 'localhost:9092'
logger_format = '%(asctime)-15s %(message)s'
logging.basicConfig(format=logger_format)
logger = logging.getLogger('data-producer')
logger.setLevel(logging.DEBUG)
value_schema = avro.load('ValueSchema.avsc')
value = {"name": "Value"}
def fetch_price(producer, symbol):
"""
helper function to retrieve stock data and send it to kafka
:param producer: instance of a kafka producer
:param symbol: symbol of the stock
:return: None
"""
logger.debug('Start to fetch stock price for %s', symbol)
try:
price = getQuotes(symbol)[0]
logger.debug('Retrieved stock info %s', price)
producer.produce(topic=topic_name, value=price)
logger.debug('Sent stock price for %s to Kafka', symbol)
except KafkaTimeoutError as timeout_error:
logger.warn('Failed to send stock price for %s to kafka, caused by: %s', (symbol, timeout_error.message))
except Exception:
logger.warn('Failed to fetch stock price for %s', symbol)
def shutdown_hook(producer):
"""
a shutdown hook to be called before the shutdown
:param producer: instance of a kafka producer
:return: None
"""
try:
logger.info('Flushing pending messages to kafka, timeout is set to 10s')
producer.flush()
logger.info('Finish flushing pending messages to kafka')
except KafkaError as kafka_error:
logger.warn('Failed to flush pending messages to kafka, caused by: %s', kafka_error.message)
finally:
try:
logger.info('Closing kafka connection')
# producer.close(10)
except Exception as e:
logger.warn('Failed to close kafka connection, caused by: %s', e.message)
if __name__ == '__main__':
# - setup command line arguments
parser = argparse.ArgumentParser()
parser.add_argument('symbol', help='the symbol of the stock to collect')
parser.add_argument('topic_name', help='the kafka topic push to')
parser.add_argument('kafka_broker', help='the location of the kafka broker')
parser.add_argument('schema_url', help='the schema registry url')
# - parse arguments
args = parser.parse_args()
symbol = args.symbol
topic_name = args.topic_name
kafka_broker = args.kafka_broker
schema_url = args.schema_url
# - instantiate avro kafka producer
producer = AvroProducer({'bootstrap.servers': kafka_broker, 'schema.registry.url': schema_url}, default_value_schema=value_schema)
# - schedule and run the fetch_price function every second
schedule.every(1).second.do(fetch_price, producer, symbol)
# - setup proper shutdown hook
atexit.register(shutdown_hook, producer)
while True:
schedule.run_pending()
time.sleep(1)