add s3 support

src/db.rs

@@ -8,12 +8,12 @@ use surrealdb::{
 use tracing::{error, instrument, trace, warn};
 use url::Url;
 
-use crate::Timer;
+use crate::{Config, Timer};
 
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct Website {
     /// The url that this data is found at
-    site: Url,
+    pub site: Url,
     /// Wether or not this link has been crawled yet
     pub crawled: bool,
     #[serde(skip_serializing)]
@@ -39,10 +39,6 @@ impl Website {
         self.crawled = true
     }
 
-    pub fn mut_url(&mut self) -> &mut Url {
-        &mut self.site
-    }
-
     #[instrument(skip_all)]
     pub async fn links_to(&self, other: Vec<Thing>, db: &Surreal<Client>) {
         let len = other.len();
@@ -149,19 +145,30 @@ pub struct Record {
     pub id: Thing,
 }
 
-pub async fn connect() -> surrealdb::Result<Surreal<Client>> {
+#[instrument(skip_all, name = "SurrealDB")]
+pub async fn connect(config: &Config<'_>) -> surrealdb::Result<Surreal<Client>> {
+    trace!("Establishing connection to surreal...");
     // Connect to the server
-    let db = Surreal::new::<Ws>("127.0.0.1:8000").await?;
+    let db = Surreal::new::<Ws>(config.surreal_url).await?;
 
+    trace!("Logging in...");
     // Signin as a namespace, database, or root user
     db.signin(Root {
-        username: "root",
-        password: "root",
+        username: config.surreal_username,
+        password: config.surreal_password,
     })
     .await?;
 
     // Select a specific namespace / database
-    db.use_ns("test").use_db("v1.2").await?;
+    db
+        .use_ns(config.surreal_ns)
+        .use_db(config.surreal_db)
+        .await?;
+    
+    let setup = include_bytes!("setup.surql");
+    let file = setup.iter().map(|c| *c as char).collect::<String>();
+
+    db.query(file).await.expect("Failed to setup surreal tables.");
 
     Ok(db)
 }

src/main.rs

@@ -6,12 +6,27 @@ use html5ever::{
     local_name, parse_document, tendril::TendrilSink, tree_builder::TreeBuilderOpts, ParseOpts,
 };
 use rcdom::RcDom;
+use s3::S3;
 use std::time::Instant;
 use surrealdb::{engine::remote::ws::Client, sql::Thing, Surreal};
 use tracing::{debug, info, instrument, trace, trace_span};
 use tracing_subscriber::EnvFilter;
 
 mod db;
+mod s3;
+
+struct Config<'a> {
+    surreal_ns: &'a str,
+    surreal_db: &'a str,
+    surreal_url: &'a str,
+    surreal_username: &'a str,
+    surreal_password: &'a str,
+
+    s3_url: &'a str,
+    s3_bucket: &'a str,
+    s3_access_key: &'a str,
+    s3_secret_key: &'a str,
+}
 
 #[tokio::main]
 async fn main() {
@@ -21,16 +36,28 @@ async fn main() {
         .without_time()
         .init();
     debug!("Starting...");
+    
+    let config = Config {
+        surreal_ns: "test",
+        surreal_db: "v1.5",
+        surreal_url: "localhost:8000",
+        surreal_username: "root",
+        surreal_password: "root",
+        s3_url: "http://localhost:9000",
+        s3_bucket: "v1.5",
+        s3_access_key: "8tUJn7e1paMFZQr0PKIT",
+        s3_secret_key: "uSMvYxNOeCejCUgXVqgTfYlUEcmiZY0xcZ91M9E0",
+    };
 
     // Would probably take these in as parameters from a cli
-    let url = "https://oliveratkinson.net/";
-    // let url = "http://localhost:5500";
-    let budget = 1000;
+    let starting_url = "https://oliveratkinson.net/";
+    let budget = 200;
     let mut crawled = 0;
 
-    let db = connect().await.expect("Failed to connect to db, aborting.");
-
-    let client = reqwest::Client::builder()
+    let s3 = S3::connect(&config).await.expect("Failed to connect to minio, aborting.");
+    let db = connect(&config).await.expect("Failed to connect to surreal, aborting.");
+    
+    let reqwest = reqwest::Client::builder()
         // .use_rustls_tls()
         .build()
         .unwrap();
@@ -40,8 +67,8 @@ async fn main() {
     let span = trace_span!("Pre-Loop");
     let pre_loop_span = span.enter();
     // Download the site
-    let mut site = Website::new(&url, false);
-    get(&mut site, &db, &client, &mut crawled).await;
+    let mut site = Website::new(&starting_url, false);
+    get(&mut site, &db, &reqwest, &s3, &mut crawled).await;
 
     drop(pre_loop_span);
 
@@ -65,7 +92,7 @@ async fn main() {
         let _ = span.enter();
 
         for mut site in uncrawled {
-            get(&mut site, &db, &client, &mut crawled).await;
+            get(&mut site, &db, &reqwest, &s3, &mut crawled).await;
             let percent = format!("{:.2}%", (crawled as f32 / budget as f32) * 100f32);
             info!("Crawled {crawled} out of {budget} pages. ({percent})");
         }
@@ -80,13 +107,14 @@ async fn main() {
 async fn get(
     site: &mut Website,
     db: &Surreal<Client>,
-    request_client: &reqwest::Client,
+    reqwest: &reqwest::Client,
+    s3: &S3,
     count: &mut usize,
 ) {
     trace!("Get: {}", site.to_string());
     let timer = Timer::start("Got page");
 
-    if let Ok(response) = request_client.get(site.to_string()).send().await {
+    if let Ok(response) = reqwest.get(site.to_string()).send().await {
         timer.stop();
 
         // Get body
@@ -98,6 +126,8 @@ async fn get(
             },
             ..Default::default()
         };
+        s3.store(&data, &site.site).await; 
+
         // Get DOM
         let dom = parse_document(RcDom::default(), opts)
             .from_utf8()
@@ -156,10 +186,10 @@ async fn walk(
                             let mut web = site.clone();
 
                             // Set url
-                            let url = web.mut_url();
+                            let mut url = web.site;
                             url.set_fragment(None); // removes #xyz
                             let joined = url.join(&attr.value).unwrap();
-                            *url = joined;
+                            web.site = joined;
 
                             // Set other attributes
                             web.crawled = false;

src/s3.rs

@@ -0,0 +1,63 @@
+use minio::s3::{
+    args::{BucketExistsArgs, MakeBucketArgs},
+    client::ClientBuilder,
+    creds::StaticProvider,
+    error::Error,
+    http::BaseUrl,
+    Client,
+};
+use tracing::{instrument, trace};
+use url::Url;
+
+use crate::Config;
+
+pub struct S3 {
+    bucket_name: String,
+    client: Client,
+}
+
+impl S3 {
+    #[instrument(skip_all, name = "S3")]
+    pub async fn connect(config: &Config<'_>) -> Result<Self, Error> {
+        let base_url = config.s3_url.parse::<BaseUrl>().unwrap();
+
+        let static_provider =
+            StaticProvider::new(&config.s3_access_key, &config.s3_secret_key, None);
+
+        let client = ClientBuilder::new(base_url)
+            .provider(Some(Box::new(static_provider)))
+            .build()?;
+
+        trace!("Checking bucket...");
+        let exists = client
+            .bucket_exists(&BucketExistsArgs::new(&config.s3_bucket).unwrap())
+            .await?;
+
+        if !exists {
+            trace!("Creating bucket...");
+            client
+                .make_bucket(&MakeBucketArgs::new(&config.s3_bucket).unwrap())
+                .await?;
+        }
+
+        trace!("Connection successfull");
+
+        Ok(Self {
+            bucket_name: config.s3_bucket.to_owned(),
+            client: client,
+        })
+    }
+
+    pub async fn store(&self, data: &str, name: &Url) {
+        if let Some(domain) = name.domain() {
+            let filename = domain.to_string() + name.path();
+
+            let _ = &self
+                .client
+                .put_object_content(&self.bucket_name, &filename, data.to_owned())
+                .send()
+                .await
+                .unwrap();
+        }
+    }
+}

src/setup.surql

@@ -0,0 +1,2 @@
+DEFINE TABLE IF NOT EXISTS website SCHEMALESS;
+DEFINE FIELD IF NOT EXISTS accessed_at ON TABLE website VALUE time::now();