tweet_sampler.py

#!/usr/bin/python3
'''
Sample from the twitter API and post results to a file or to Kafka.
To use, set credientials as enviornment variables, e.g.
export TWITTER_ACCESS_TOKEN=... 
or source myfile
where myfile exports the authorization variables 
'''
import twython, json, re, argparse, subprocess, os, sys, time
from socket import timeout
from kafka import SimpleProducer, KafkaClient

####################
#    Constants
####################
access_token = os.environ["TWITTER_ACCESS_TOKEN"]
access_token_secret = os.environ["TWITTER_SECRET_TOKEN"] 
consumer_key = os.environ["TWITTER_CONSUMER_TOKEN"]
consumer_secret = os.environ["TWITTER_CONSUMER_SECRET"]

class StrawStreamer(twython.TwythonStreamer):

    def __init__(self, APP_KEY, APP_SECRET, OAUTH_TOKEN, OAUTH_TOKEN_SECRET, outfile):
        super(StrawStreamer, self).__init__(APP_KEY, APP_SECRET,OAUTH_TOKEN, OAUTH_TOKEN_SECRET)
        self.outfile=outfile

    def on_success(self, data):
        if 'text' in data:
            self.outfile.write((json.dumps(data)+u'\n').encode('utf-8'))

    def on_error(self, status_code, data):
        print(status_code)

class KafkaStrawStreamer(twython.TwythonStreamer):
    def __init__(self, APP_KEY, APP_SECRET, OAUTH_TOKEN, OAUTH_TOKEN_SECRET, host, port):
        super(KafkaStrawStreamer, self).__init__(APP_KEY, APP_SECRET,OAUTH_TOKEN, OAUTH_TOKEN_SECRET)

        # connect to Kafka
        print("Connecting to Kafka node {0}:{1}".format(host, port))
        kafka = KafkaClient("{0}:{1}".format(host, port))
        self.producer = BufferedSimpleProducer(kafka, 100)

    def on_success(self, data):
        # TODO: add message queue so we can pass messages in bulk
        if 'text' in data:
            msg = (json.dumps(data)+u'\n').encode('utf-8')
            self.producer.send_messages(args.topic, msg)

    def on_error(self, status_code, data):
        print(status_code)    

class BufferedSimpleProducer:
    def __init__(self, kafka, chunk_size):
        self.producer = SimpleProducer(kafka)
        self.queues = {}
        self.chunk_size = chunk_size

    def send_messages(self, topic, msg):
        if topic not in self.queues:
            self.queues[topic]=[]
        if len(self.queues[topic])<self.chunk_size:
            self.queues[topic].append(msg)
        else:
            self.producer.send_messages(topic, *(self.queues[topic]))
            print("Sent {0} documents to Kafka.".format(len(self.queues[topic])))
            self.queues[topic] = []

if __name__=="__main__":

    # arg parsing
    parser = argparse.ArgumentParser(description="Python twitter firehose sampler")
    parser.add_argument("-f","--file", help="Output will be appended to this file.")
    parser.add_argument("-k","--kafka", help="A kafka broker node")
    parser.add_argument("-t","--topic", help="A kafka topic")
    parser.add_argument("-p","--port", default=9092, help="A kafka port, default 9200")
    args = parser.parse_args()
    if args.file is None and args.kafka is None:
        raise RuntimeError("Need either an output file or a kafka host")
    if args.kafka is not None and args.topic is None:
        raise RuntimeError("Need a topic for Kafka")

    # connect to twitter API
    twitter = twython.Twython(consumer_key, consumer_secret)
    
    # write to file or to Kafka
    if args.file is not None:
        with open(args.file, "ab") as f:
            while True:
                try:
                    stream = StrawStreamer(consumer_key, consumer_secret, access_token, access_token_secret, f)
                    stream.statuses.sample(language="en")
                except timeout as e:
                    print("GOT SOCKET ERROR: {0}".format(e))
                    print("Retrying connection after 500 second wait...")
                    f.flush()
                    time.sleep(500)
    else:
        while True:
            try:
                stream = KafkaStrawStreamer(consumer_key, consumer_secret, access_token, access_token_secret, args.kafka, args.port)
                stream.statuses.sample(language="en")
            except timeout as e:
                print("GOT SOCKET ERROR: {0}".format(e))
                print("Retrying connection after 500 second wait...")
                f.flush()
                time.sleep(500)