Adapting example from egg crate + twitch-firehose (WIP)

Cargo.toml

@@ -0,0 +1,32 @@
+[package]
+name = "twitter-firehose"
+version = "0.0.0"
+edition = "2018"
+authors = [
+  "Brandon Thomas <bt@brand.io>",
+  "Brandon Thomas <echelon@gmail.com>"
+]
+
+[[bin]]
+name = "twitter-firehose"
+path = "src/main.rs"
+
+[dependencies]
+anyhow = "1.0.38"
+easyenv = "0.0.4"
+log = "0.4.14"
+prost = "0.7.0"
+rustls = "0.18"
+serde = "1.0.123"
+serde_derive = "1.0.123"
+tokio = { version = "0.2", features = ["full"] }
+toml = "0.5.8"
+
+redis = { version = "0.17.0", features = ["tls", "async-std-tls-comp", "tokio-tls-comp", "connection-manager"] }
+
+# Twitter
+egg-mode = "0.15.0"
+
+[build-dependencies]
+anyhow = "1.0.38"
+prost-build = "0.7.0"

src/clients/mod.rs

@@ -0,0 +1 @@
+pub mod redis_client;

src/clients/redis_client.rs

@@ -0,0 +1,65 @@
+use anyhow::anyhow;
+use crate::AnyhowResult;
+use crate::secrets::RedisSecrets;
+use log::{debug, info};
+use redis::AsyncCommands;
+use redis::aio::{Connection, ConnectionManager};
+use redis;
+use std::thread;
+use std::time::Duration;
+
+pub struct RedisClient {
+  secrets: RedisSecrets,
+
+  /// This manages our connection.
+  /// It's meant to be absent when not connected.
+  connection_manager: Option<ConnectionManager>,
+
+  /// Attempts to retry a Redis command.
+  max_retry_count: u32,
+}
+
+impl RedisClient {
+  pub fn new(secrets: &RedisSecrets, retry_count: u32) -> Self {
+    Self {
+      secrets: secrets.clone(),
+      connection_manager: None,
+      max_retry_count: retry_count,
+    }
+  }
+
+  pub async fn connect(&mut self) -> AnyhowResult<()> {
+    let client = redis::Client::open(self.secrets.connection_url())?;
+
+    let connection_manager = client.get_tokio_connection_manager().await?;
+    self.connection_manager = Some(connection_manager);
+
+    Ok(())
+  }
+
+  pub async fn publish_bytes(&mut self, channel: &str, message: &[u8]) -> AnyhowResult<u32> {
+    for i in 0 .. self.max_retry_count {
+      match self.try_publish_bytes(channel, message).await {
+        Ok(r) => return Ok(r),
+        Err(_) => {
+          debug!("Redis publish_bytes failed; retrying (attempt {})", i);
+          continue;
+        },
+      }
+    }
+    Err(anyhow!("Retry count elapsed."))
+  }
+
+  async fn try_publish_bytes(&mut self, channel: &str, message: &[u8]) -> AnyhowResult<u32> {
+    let connection_manager = match &mut self.connection_manager {
+      None => {
+        self.connect().await?;
+        return Err(anyhow!("Not connected"));
+      },
+      Some(cm) => cm,
+    };
+
+    let result = connection_manager.publish(channel, message).await?;
+    Ok(result)
+  }
+}

src/main.rs

@@ -0,0 +1,76 @@
+#![allow(dead_code)]
+#![allow(unused_mut)]
+#![allow(unused_imports)]
+#![warn(unused_must_use)]
+//#![allow(warnings)]
+
+#[macro_use]
+extern crate serde_derive;
+
+mod clients;
+mod secrets;
+//mod util;
+
+//use futures::TryStreamExt;
+use egg_mode::stream::StreamMessage;
+use crate::secrets::Secrets;
+use crate::clients::redis_client::RedisClient;
+
+pub type AnyhowResult<T> = anyhow::Result<T>;
+
+const ENV_PUBLISH_TOPIC : &'static str = "PUBLISH_TOPIC";
+const ENV_PUBLISH_TOPIC_DEFAULT : &'static str = "firehose";
+
+const ENV_REDIS_MAX_RETRY_COUNT : &'static str = "REDIS_MAX_RETRY_COUNT";
+const ENV_REDIS_MAX_RETRY_COUNT_DEFAULT : u32 = 3;
+
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+
+  easyenv::init_env_logger(None);
+
+  let redis_publish_topic = easyenv::get_env_string_or_default(
+    ENV_PUBLISH_TOPIC, ENV_PUBLISH_TOPIC_DEFAULT);
+
+  let redis_max_retry_count = easyenv::get_env_num(
+    ENV_REDIS_MAX_RETRY_COUNT,
+    ENV_REDIS_MAX_RETRY_COUNT_DEFAULT)?;
+
+  let secrets = Secrets::from_file("secrets.toml")?;
+
+  let mut redis_client = RedisClient::new(
+    &secrets.redis,
+    redis_max_retry_count
+  );
+
+  redis_client.connect().await?;
+
+  // We can capture: Mentions, Retweets, Replies.
+  // We cannot capture: Follows.
+
+  /*let config = common::Config::load().await;
+  println!("Streaming tweets containing popular programming languages (and also Rust)");
+  println!("Ctrl-C to quit\n");
+
+  let stream = egg_mode::stream::filter()
+      .follow(&[1297106371238932481]) // vocodes
+      //.track(&["rustlang", "python", "java", "javascript"])
+      .language(&["en"])
+      .start(&config.token)
+      .try_for_each(|m| {
+          if let StreamMessage::Tweet(tweet) = m {
+              common::print_tweet(&tweet);
+              println!("──────────────────────────────────────");
+          } else {
+              println!("{:?}", m);
+          }
+          futures::future::ok(())
+      });
+  if let Err(e) = stream.await {
+      println!("Stream error: {}", e);
+      println!("Disconnected")
+  }
+*/
+
+  Ok(())
+}

src/secrets.rs

@@ -0,0 +1,49 @@
+use crate::AnyhowResult;
+use std::io::Read;
+
+#[derive(Deserialize)]
+pub struct Secrets {
+  pub redis: RedisSecrets,
+  pub twitter: TwitterSecrets,
+}
+
+#[derive(Deserialize, Clone)]
+pub struct RedisSecrets {
+  pub username: String,
+  pub password: String,
+  pub host: String,
+  pub port: u32,
+}
+
+#[derive(Deserialize, Clone)]
+pub struct TwitterSecrets {
+  pub api_key: String,
+  pub api_secret_key: String,
+}
+
+impl Secrets {
+  pub fn from_file(filename: &str) -> AnyhowResult<Self> {
+    let mut file = std::fs::File::open(filename)?;
+    let mut buffer = String::new();
+    file.read_to_string(&mut buffer)?;
+    let secrets = toml::from_str(&buffer)?;
+    Ok(secrets)
+  }
+}
+
+impl RedisSecrets {
+  pub fn connection_url(&self) -> String {
+    format!("rediss://{}:{}@{}:{}", self.username, self.password, self.host, self.port)
+  }
+}
+
+/*impl TwitterSecrets {
+  pub fn get_user_config(&self) -> AnyhowResult<UserConfig> {
+    let config = UserConfig::builder()
+      .name(self.username.to_string())
+      .token(self.oauth_access_token.to_string())
+      .enable_all_capabilities()
+      .build()?;
+    Ok(config)
+  }
+}*/

src/util.rs

@@ -0,0 +1,184 @@
+
+pub struct Config {
+  pub token: egg_mode::Token,
+  pub user_id: u64,
+  pub screen_name: String,
+}
+
+impl Config {
+  pub async fn load() -> Self {
+    let a1 = Config::load_inner().await;
+    if let Some(conf) = a1 {
+      return conf;
+    }
+
+    Config::load_inner().await.unwrap()
+  }
+
+  /// This needs to be a separate function so we can retry after creating the
+  /// twitter_settings file. Idealy we would recurse, but that requires boxing
+  /// the output which doesn't seem worthwhile
+  async fn load_inner() -> Option<Self> {
+    //IMPORTANT: make an app for yourself at apps.twitter.com and get your
+    //key/secret into these files; these examples won't work without them
+    let consumer_key = include_str!("consumer_key").trim();
+    let consumer_secret = include_str!("consumer_secret").trim();
+
+    let con_token = egg_mode::KeyPair::new(consumer_key, consumer_secret);
+
+    let mut config = String::new();
+    let user_id: u64;
+    let username: String;
+    let token: egg_mode::Token;
+
+    //look at all this unwrapping! who told you it was my birthday?
+    if let Ok(mut f) = std::fs::File::open("twitter_settings") {
+      f.read_to_string(&mut config).unwrap();
+
+      let mut iter = config.split('\n');
+
+      username = iter.next().unwrap().to_string();
+      user_id = u64::from_str_radix(&iter.next().unwrap(), 10).unwrap();
+      let access_token = egg_mode::KeyPair::new(
+        iter.next().unwrap().to_string(),
+        iter.next().unwrap().to_string(),
+      );
+      token = egg_mode::Token::Access {
+        consumer: con_token,
+        access: access_token,
+      };
+
+      if let Err(err) = egg_mode::auth::verify_tokens(&token).await {
+        println!("We've hit an error using your old tokens: {:?}", err);
+        println!("We'll have to reauthenticate before continuing.");
+        std::fs::remove_file("twitter_settings").unwrap();
+      } else {
+        println!("Welcome back, {}!\n", username);
+      }
+    } else {
+      let request_token = egg_mode::auth::request_token(&con_token, "oob").await.unwrap();
+
+      println!("Go to the following URL, sign in, and give me the PIN that comes back:");
+      println!("{}", egg_mode::auth::authorize_url(&request_token));
+
+      let mut pin = String::new();
+      std::io::stdin().read_line(&mut pin).unwrap();
+      println!("");
+
+      let tok_result = egg_mode::auth::access_token(con_token, &request_token, pin)
+        .await
+        .unwrap();
+
+      token = tok_result.0;
+      user_id = tok_result.1;
+      username = tok_result.2;
+
+      match token {
+        egg_mode::Token::Access {
+          access: ref access_token,
+          ..
+        } => {
+          config.push_str(&username);
+          config.push('\n');
+          config.push_str(&format!("{}", user_id));
+          config.push('\n');
+          config.push_str(&access_token.key);
+          config.push('\n');
+          config.push_str(&access_token.secret);
+        }
+        _ => unreachable!(),
+      }
+
+      let mut f = std::fs::File::create("twitter_settings").unwrap();
+      f.write_all(config.as_bytes()).unwrap();
+
+      println!("Welcome, {}, let's get this show on the road!", username);
+    }
+
+    //TODO: Is there a better way to query whether a file exists?
+    if std::fs::metadata("twitter_settings").is_ok() {
+      Some(Config {
+        token: token,
+        user_id: user_id,
+        screen_name: username,
+      })
+    } else {
+      None
+    }
+  }
+}
+
+
+pub fn print_tweet(tweet: &egg_mode::tweet::Tweet) {
+  if let Some(ref user) = tweet.user {
+    println!(
+      "{} (@{}) posted at {}",
+      Paint::blue(&user.name),
+      Paint::bold(Paint::blue(&user.screen_name)),
+      tweet.created_at.with_timezone(&chrono::Local)
+    );
+  }
+
+  if let Some(ref screen_name) = tweet.in_reply_to_screen_name {
+    println!("➜ in reply to @{}", Paint::blue(screen_name));
+  }
+
+  if let Some(ref status) = tweet.retweeted_status {
+    println!("{}", Paint::red("Retweet ➜"));
+    print_tweet(status);
+    return;
+  } else {
+    println!("{}", Paint::green(&tweet.text));
+  }
+
+  if let Some(source) = &tweet.source {
+    println!("➜ via {} ({})", source.name, source.url);
+  }
+
+  if let Some(ref place) = tweet.place {
+    println!("➜ from: {}", place.full_name);
+  }
+
+  if let Some(ref status) = tweet.quoted_status {
+    println!("{}", Paint::red("➜ Quoting the following status:"));
+    print_tweet(status);
+  }
+
+  if !tweet.entities.hashtags.is_empty() {
+    println!("➜ Hashtags contained in the tweet:");
+    for tag in &tweet.entities.hashtags {
+      println!("  {}", tag.text);
+    }
+  }
+
+  if !tweet.entities.symbols.is_empty() {
+    println!("➜ Symbols contained in the tweet:");
+    for tag in &tweet.entities.symbols {
+      println!("  {}", tag.text);
+    }
+  }
+
+  if !tweet.entities.urls.is_empty() {
+    println!("➜ URLs contained in the tweet:");
+    for url in &tweet.entities.urls {
+      if let Some(expanded_url) = &url.expanded_url {
+        println!("  {}", expanded_url);
+      }
+    }
+  }
+
+  if !tweet.entities.user_mentions.is_empty() {
+    println!("➜ Users mentioned in the tweet:");
+    for user in &tweet.entities.user_mentions {
+      println!("  {}", Paint::bold(Paint::blue(&user.screen_name)));
+    }
+  }
+
+  if let Some(ref media) = tweet.extended_entities {
+    println!("➜ Media attached to the tweet:");
+    for info in &media.media {
+      println!("  A {:?}", info.media_type);
+    }
+  }
+}
+