Compare commits

11 Commits

Author SHA1 Message Date
4989a59ddf checkpoint 2025-07-10 18:46:25 -06:00
6fc71c7a78 add speed improvements 2025-03-21 12:14:29 -06:00
96a3ca092a :) 2025-03-21 12:11:05 -06:00
b750d88d48 working filesystem storage 2025-03-21 11:42:43 -06:00
808790a7c3 file patch; 2025-03-21 07:11:51 +00:00
2de01b2a0e remove removed code 2025-03-21 06:48:39 +00:00
be0fd5505b i think the files work better 2025-03-21 06:48:17 +00:00
a23429104c dead code removal 2025-03-21 06:03:34 +00:00
66581cc453 getting there 2025-03-21 05:59:40 +00:00
7df19a480f updates 2025-03-20 15:11:01 -06:00
b9c1f0b492 readme updates 2025-03-19 15:05:32 -06:00
13 changed files with 268 additions and 2641 deletions

1
.gitignore vendored
View File

@@ -5,3 +5,4 @@ perf.data
flamegraph.svg
perf.data.old
/docker/logs/*
/downloaded

2311
Cargo.lock generated

File diff suppressed because it is too large Load Diff

View File

@@ -1,5 +1,5 @@
[package]
name = "surreal_spider"
name = "internet_mapper"
version = "0.1.0"
edition = "2021"
@@ -11,8 +11,8 @@ metrics-exporter-prometheus = { version = "0.16.2", features=["http-listener"]}
# minio = "0.1.0"
minio = {git="https://github.com/minio/minio-rs.git", rev = "c28f576"}
reqwest = { version = "0.12", features = ["gzip", "default", "rustls-tls"] }
rusqlite = { version = "0.34.0", features = ["bundled"] }
serde = { version = "1.0", features = ["derive"] }
surrealdb = "2.2"
tokio = { version="1.41.0", features = ["full"] }
toml = "0.8.20"
tracing = "0.1"

View File

@@ -3,13 +3,7 @@ surreal_url = "localhost:8000"
surreal_username = "root"
surreal_password = "root"
surreal_ns = "test"
surreal_db = "v1.15.4"
# Minio config
s3_bucket = "v1.15.4"
s3_url = "http://localhost:9000"
s3_access_key = "3ptjsHhRHCHlpCmgFy9n"
s3_secret_key = "68CmV07YExeCxb8kJhosSauEizj5CAE7PINZIfQz"
surreal_db = "v1.19.2"
# Crawler config
crawl_filter = "en.wikipedia.com"

View File

@@ -2,11 +2,24 @@
Crawls sites saving all the found links to a surrealdb database. It then proceeds to take batches of 100 uncrawled links untill the crawl budget is reached. It saves the data of each site in a minio database.
### TODO
- [ ] Domain filtering - prevent the crawler from going on alternate versions of wikipedia.
- [ ] Conditionally save content - based on filename or file contents
- [ ] GUI / TUI ?
- [ ] Better asynchronous getting of the sites. Currently it all happens serially.3/19/25: Took 20min to crawl 100 pages
This ment we stored 100 pages, 142,997 urls, and 1,425,798 links between the two.
- [x] GUI / TUI ? - Graphana
- [x] Better asynchronous getting of the sites. Currently it all happens serially.
- [ ] Allow for storing asynchronously
3/17/25: Took >1hr to crawl 100 pages
3/19/25: Took 20min to crawl 1000 pages
This ment we stored 1000 pages, 142,997 urls, and 1,425,798 links between the two.
3/20/25: Took 5min to crawl 1000 pages
3/21/25: Took 3min to crawl 1000 pages
# About
![Screenshot](/pngs/graphana.png)

View File

@@ -14,22 +14,6 @@ services:
- --pass
- root
- rocksdb:/mydata/database.db
minio:
image: quay.io/minio/minio
ports:
- 9000:9000
- 9001:9001
environment:
- MINIO_ROOT_USER=root
- MINIO_ROOT_PASSWORD=an8charpassword
- MINIO_PROMETHEUS_AUTH_TYPE=public
volumes:
- minio_storage:/data
command:
- server
- /data
- --console-address
- ":9001"
alloy:
image: grafana/alloy:latest

View File

@@ -8,13 +8,10 @@ scrape_configs:
# change this your machine's ip, localhost won't work
# because localhost refers to the docker container.
- targets: ['172.20.239.48:2500']
#- targets: ['192.168.8.209:2500']
- job_name: loki
static_configs:
- targets: ['loki:3100']
- job_name: prometheus
static_configs:
- targets: ['localhost:9090']
- job_name: minio
metrics_path: /minio/v2/metrics/cluster
static_configs:
- targets: ['minio:9000']

BIN
pngs/graphana.png Normal file

Binary file not shown.

After

Width:  |  Height:  |  Size: 264 KiB

165
src/db.rs
View File

@@ -1,35 +1,15 @@
use base64::{
alphabet,
engine::{self, general_purpose},
Engine,
};
use metrics::counter;
use serde::{ser::SerializeStruct, Deserialize, Serialize};
use std::{fmt::Debug, sync::LazyLock, time::Instant};
use surrealdb::{
engine::remote::ws::{Client, Ws},
opt::auth::Root,
sql::Thing,
Error::Api,
Response, Surreal,
};
use tokio::sync::Mutex;
use tracing::{error, instrument, trace, warn};
use rusqlite::Connection;
use std::fmt::Debug;
use serde::{Deserialize, Serialize};
use tracing::{error, instrument, trace};
use url::Url;
use crate::{Config, Timer};
use crate::Config;
// static LOCK: LazyLock<Arc<Mutex<bool>>> = LazyLock::new(|| Arc::new(Mutex::new(true)));
static LOCK: LazyLock<Mutex<bool>> = LazyLock::new(|| Mutex::new(true));
const STORE: &str = "surql_store_calls";
const CUSTOM_ENGINE: engine::GeneralPurpose =
engine::GeneralPurpose::new(&alphabet::URL_SAFE, general_purpose::NO_PAD);
const TIME_SPENT_ON_LOCK: &'static str = "surql_lock_waiting_ms";
const STORE: &'static str = "surql_store_calls";
const LINK: &'static str = "surql_link_calls";
#[derive(Deserialize, Clone)]
#[derive(Serialize, Deserialize, Clone, Eq, PartialEq, Hash)]
pub struct Website {
/// The url that this data is found at
pub site: Url,
@@ -37,18 +17,6 @@ pub struct Website {
pub crawled: bool,
}
impl Serialize for Website {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer {
let mut state = serializer.serialize_struct("Website", 2)?;
state.serialize_field("crawled", &self.crawled)?;
// to_string() calls the correct naming of site
state.serialize_field("site", &self.site.to_string())?;
state.end()
}
}
// manual impl to make tracing look nicer
impl Debug for Website {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -74,84 +42,17 @@ impl Website {
self.crawled = true
}
pub fn get_url_as_string(site: &Url) -> String {
let domain = match site.domain() {
Some(s) => s.to_string(),
None => {
warn!("Failed to get domain of URL: {}, falling back to 'localhost'", site.to_string());
"localhost".to_string()
}
};
let path = site.path();
domain + path
}
pub fn get_url_as_b64_path(site: &Url) -> String {
let domain = site.domain().unwrap_or("DOMAIN").to_string();
let path = &CUSTOM_ENGINE.encode(site.path());
domain + path
}
#[instrument(skip_all)]
pub async fn links_to(&self, other: Vec<Thing>, db: &Surreal<Client>) {
let len = other.len();
if len == 0 {
return;
}
let from = &self.site;
// let to = other.site.to_string();
trace!("Linking {} pages to {from}", other.len());
let msg = format!("Linked {len} pages to {from}");
let timer = Timer::start(&msg);
// prevent the timer from being dropped instantly.
let _ = timer;
counter!(LINK).increment(1);
match db
.query("COUNT(RELATE (SELECT id FROM website WHERE site = $in) -> links_to -> $out)")
.bind(("in", from.clone()))
.bind(("out", other))
.await
{
Ok(mut e) => {
// The relate could technically "fail" (not relate anything), this just means that
// the query was ok.
let _: Response = e;
if let Ok(vec) = e.take(0) {
let _: Vec<usize> = vec;
if let Some(num) = vec.get(0) {
if *num == len {
trace!("Link for {from} OK - {num}/{len}");
return;
} else {
error!("Didn't link all the records. {num}/{len}. Surreal response: {:?}", e);
return;
}
}
}
error!("Linking request succeeded but couldn't verify the results.");
}
Err(e) => {
error!("{}", e.to_string());
}
}
}
// Insert ever item in the vec into surreal, crawled state will be preserved as TRUE
// if already in the database as such or incoming data is TRUE.
pub async fn store_all(all: Vec<Self>, db: &Surreal<Client>) -> Vec<Thing> {
pub async fn store_all(all: Vec<Self>, db: &Connection) {
counter!(STORE).increment(1);
let mut things = Vec::with_capacity(all.len());
// TODO this only allows for one thread to be in the database at a time.
// This is currently required since otherwise we get write errors.
// If the default `crawled` is set to false, we might not need to write any more
// than just the name. `accessed_at` is fun but not needed.
let now = Instant::now();
let lock = LOCK.lock().await;
counter!(TIME_SPENT_ON_LOCK).increment(now.elapsed().as_millis() as u64);
rusqlite::ParamsFromIter;
db.execute("",
params![]
);
match db
.query(
@@ -167,19 +68,12 @@ impl Website {
{
Ok(mut id) => match id.take::<Vec<Thing>>(0) {
Ok(mut x) => things.append(&mut x),
Err(err) => match err {
Api(error) => {
eprintln!("{:?}", error);
error!("{:?}", error);
}
_ => error!("{:?}", err),
},
Err(err) => error!("{:?}", err),
},
Err(err) => {
error!("{:?}", err);
}
}
drop(lock);
things
}
}
@@ -196,31 +90,10 @@ pub struct Record {
pub id: Thing,
}
#[instrument(skip_all, name = "SurrealDB")]
pub async fn connect(config: &Config) -> surrealdb::Result<Surreal<Client>> {
trace!("Establishing connection to surreal...");
#[instrument(skip_all, name = "sqlite_connect")]
pub async fn connect(config: &Config) -> Result<Connection, rusqlite::Error> {
trace!("Establishing connection to sqlite...");
// Connect to the server
let db = Surreal::new::<Ws>(&config.surreal_url).await?;
trace!("Logging in...");
// Signin as a namespace, database, or root user
db.signin(Root {
username: &config.surreal_username,
password: &config.surreal_password,
})
.await?;
// Select a specific namespace / database
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)
Connection::open("./squeelite.db")
}

70
src/filesystem.rs Normal file
View File

@@ -0,0 +1,70 @@
use std::{ffi::OsStr, path::PathBuf};
use tokio::fs;
use tracing::{debug, error, instrument, trace, warn};
use url::Url;
#[instrument(skip(data))]
pub async fn store(data: &str, url: &Url) {
// extract data from url to save it accurately
let url_path = PathBuf::from("./downloaded/".to_string() + url.domain().unwrap_or("UnknownDomain") + url.path());
// if it's a file
let (basepath, filename) = if url_path.extension().filter(valid_file_extension).is_some() {
// get everything up till the file
let basepath = url_path.ancestors().skip(1).take(1).collect::<PathBuf>();
// get the file name
let filename = url_path.file_name().expect("This should exist").to_string_lossy();
trace!("Save path: {:?} and base path: {:?}", &url_path, &basepath);
(basepath, filename.to_string())
} else {
(url_path.clone(), "index.html".into())
};
debug!("Writing at: {:?} {:?}", basepath, filename);
// create the folders
if let Err(err) = fs::create_dir_all(&basepath).await {
error!("Dir creation: {err} {:?}", basepath);
} else {
// FIXME I don't think this handles index.html files well...
// TODO this should probably append .html to non-described files
// create the file if that was successful
if let Err(err) = fs::write(&basepath.join(filename), data).await {
error!("File creation: {err} {:?}", url_path);
}
}
}
fn valid_file_extension(take: &&OsStr) -> bool {
let los = take.to_string_lossy();
let all = los.split('.');
match all.last() {
Some(s) => {
match s.to_lowercase().as_str() {
"html" => true,
"css" => true,
"js" => true,
"ts" => true,
"otf" => true, // font
"png" => true,
"svg" => true,
"jpg" => true,
"jpeg" => true,
"mp4" => true,
"mp3" => true,
"webp" => true,
"pdf" => true,
"json" => true,
"xml" => true,
_ => {
warn!("Might be forgetting a file extension: {s}");
false
}
}
},
None => false,
}
}

View File

@@ -3,30 +3,26 @@
extern crate html5ever;
use std::{
fs::File,
io::Read,
net::{IpAddr, Ipv4Addr},
time::Instant,
collections::HashSet, fs::File, io::Read, net::{IpAddr, Ipv4Addr}
};
use db::{connect, Website};
use metrics::{counter, gauge};
use metrics_exporter_prometheus::PrometheusBuilder;
use s3::S3;
use serde::Deserialize;
use surrealdb::{engine::remote::ws::Client, Surreal};
use tokio::task::JoinSet;
use tracing::{debug, error, info, instrument, trace, trace_span, warn};
use tracing::{debug, error, info, instrument, level_filters::LevelFilter, trace, trace_span};
use tracing_subscriber::{fmt, layer::SubscriberExt, EnvFilter, Layer, Registry};
mod db;
mod parser;
mod s3;
mod filesystem;
const GET_METRIC: &'static str = "total_gets";
const GET_IN_FLIGHT: &'static str = "gets_in_flight";
const SITES_CRAWLED: &'static str = "pages_crawled";
const BEING_PROCESSED: &'static str = "pages_being_processed";
const GET_METRIC: &str = "total_gets";
const GET_IN_FLIGHT: &str = "gets_in_flight";
const SITES_CRAWLED: &str = "pages_crawled";
const BEING_PROCESSED: &str = "pages_being_processed";
#[derive(Deserialize)]
struct Config {
@@ -36,26 +32,21 @@ struct Config {
surreal_username: String,
surreal_password: String,
s3_url: String,
s3_bucket: String,
s3_access_key: String,
s3_secret_key: String,
crawl_filter: String,
budget: usize,
}
#[tokio::main]
async fn main() {
let total_runtime = Timer::start("Completed");
let writer = std::fs::OpenOptions::new()
.append(true)
.create(true)
.open("./docker/logs/tracing.log")
.expect("Couldn't make log file!");
let filter = EnvFilter::from_default_env();
let filter = EnvFilter::builder()
.with_default_directive(LevelFilter::DEBUG.into())
.from_env_lossy();
let registry = Registry::default().with(
fmt::layer()
@@ -78,7 +69,7 @@ async fn main() {
.install()
.expect("failed to install recorder/exporter");
debug!("Starting...");
info!("Starting...");
// Would probably take these in as parameters from a cli
let starting_url = "https://en.wikipedia.org/";
// When getting uncrawled pages, name must contain this variable. "" will effectively get ignored.
@@ -95,9 +86,6 @@ async fn main() {
let db = connect(&config)
.await
.expect("Failed to connect to surreal, aborting.");
let s3 = S3::connect(&config)
.await
.expect("Failed to connect to minio, aborting.\n\nThis probably means you need to login to the minio console and get a new access key!\n\n(Probably here) http://localhost:9001/access-keys/new-account\n\n");
let reqwest = reqwest::Client::builder()
// .use_rustls_tls()
@@ -110,8 +98,8 @@ async fn main() {
let span = trace_span!("Pre-Loop");
let pre_loop_span = span.enter();
// Download the site
let site = Website::new(&starting_url, false);
process(site, db.clone(), reqwest.clone(), s3.clone()).await;
let site = Website::new(starting_url, false);
process(site, db.clone(), reqwest.clone()).await;
drop(pre_loop_span);
@@ -125,28 +113,23 @@ async fn main() {
};
let uncrawled = get_uncrawled_links(&db, get_num, config.crawl_filter.clone()).await;
if uncrawled.len() == 0 {
if uncrawled.is_empty() {
info!("Had more budget but finished crawling everything.");
return;
}
debug!("Crawling {} pages...", uncrawled.len());
let span = trace_span!("Crawling");
let _ = span.enter();
{
let mut futures = JoinSet::new();
for site in uncrawled {
gauge!(BEING_PROCESSED).increment(1);
futures.spawn(process(site, db.clone(), reqwest.clone(), s3.clone()));
futures.spawn(process(site, db.clone(), reqwest.clone()));
// let percent = format!("{:.2}%", (crawled as f32 / budget as f32) * 100f32);
// info!("Crawled {crawled} out of {budget} pages. ({percent})");
}
debug!("Joining {} futures...", futures.len());
let c = counter!(SITES_CRAWLED);
// As futures complete runs code in while block
while let Some(_) = futures.join_next().await {
while futures.join_next().await.is_some() {
c.increment(1);
gauge!(BEING_PROCESSED).decrement(1);
crawled += 1;
@@ -155,33 +138,27 @@ async fn main() {
}
drop(span);
debug!("Done");
drop(total_runtime);
info!("Done");
}
#[instrument(skip(db, s3, reqwest))]
#[instrument(skip(db, reqwest))]
/// Downloads and crawls and stores a webpage.
/// It is acceptable to clone `db`, `reqwest`, and `s3` because they all use `Arc`s internally. - Noted by Oliver
async fn process(mut site: Website, db: Surreal<Client>, reqwest: reqwest::Client, s3: S3) {
async fn process(mut site: Website, db: Surreal<Client>, reqwest: reqwest::Client) {
// METRICS
trace!("Process: {}", &site.site);
let timer = Timer::start("Built request");
// Build the request
let request_builder = reqwest.get(&site.site.to_string());
// METRICS
timer.stop();
let request_builder = reqwest.get(site.site.to_string());
// METRICS
let g = gauge!(GET_IN_FLIGHT);
g.increment(1);
let timer = Timer::start("Got page");
// Send the http request (get)
if let Ok(response) = request_builder.send().await {
// METRICS
timer.stop();
g.decrement(1);
counter!(GET_METRIC).increment(1);
@@ -190,22 +167,30 @@ async fn process(mut site: Website, db: Surreal<Client>, reqwest: reqwest::Clien
.text()
.await
.expect("Failed to read http response's body!");
// Store document
s3.store(&data, &site.site).await;
filesystem::store(&data, &site.site).await;
// Parse document and get relationships
let sites = parser::parse(&site, &data).await;
// update self in db
site.set_crawled();
Website::store_all(vec![site.clone()], &db).await;
Website::store_all(vec![site], &db).await;
// De-duplicate this list
let prev_len = sites.len();
let set = sites.into_iter().fold(HashSet::new(), |mut set,item| {
set.insert(item);
set
});
let de_dupe_sites: Vec<Website> = set.into_iter().collect();
let diff = prev_len - de_dupe_sites.len();
trace!("Saved {diff} from being entered into the db by de-duping");
// Store all the other sites so that we can link to them.
// let mut links_to = Vec::new();
let others = Website::store_all(sites, &db).await;
let _ = Website::store_all(de_dupe_sites, &db).await;
// Make the database's links reflect the html links between sites
site.links_to(others, &db).await;
} else {
error!("Failed to get: {}", &site.site);
}
@@ -234,33 +219,3 @@ async fn get_uncrawled_links(
.expect("Returned websites couldn't be parsed")
}
pub struct Timer<'a> {
start: Instant,
msg: &'a str,
}
impl<'a> Timer<'a> {
#[inline]
pub fn start(msg: &'a str) -> Self {
Self {
start: Instant::now(),
msg,
}
}
pub fn stop(&self) -> f64 {
let dif = self.start.elapsed().as_micros();
let ms = dif as f64 / 1000.;
if ms > 200. {
warn!("{}", format!("{} in {:.3}ms", self.msg, ms));
}
ms
}
}
impl Drop for Timer<'_> {
fn drop(&mut self) {
self.stop();
}
}

View File

@@ -5,19 +5,21 @@ use html5ever::tokenizer::{BufferQueue, TokenizerResult};
use html5ever::tokenizer::{StartTag, TagToken};
use html5ever::tokenizer::{Token, TokenSink, TokenSinkResult, Tokenizer, TokenizerOpts};
use html5ever::{local_name, tendril::*};
use tracing::instrument;
use tracing::{debug, error, instrument, trace, warn};
use url::Url;
use crate::db::Website;
use crate::Timer;
impl TokenSink for Website {
type Handle = Vec<Website>;
#[instrument(skip(token, _line_number))]
fn process_token(&self, token: Token, _line_number: u64) -> TokenSinkResult<Self::Handle> {
match token {
TagToken(tag) => {
if tag.kind == StartTag {
match tag.name {
// this should be all the html elements that have links
local_name!("a")
| local_name!("audio")
| local_name!("area")
@@ -32,23 +34,18 @@ impl TokenSink for Website {
let attr_name = attr.name.local.to_string();
if attr_name == "src" || attr_name == "href" || attr_name == "data"
{
// Get clone of the current site object
let mut web = self.clone();
trace!("Found `{}` in html `{}` tag", &attr.value, tag.name);
let url = try_get_url(&self.site, &attr.value);
// Set url
let mut url = web.site;
url.set_fragment(None); // removes #xyz
let joined = url
.join(&attr.value)
.expect("Failed to join url during parsing!");
web.site = joined;
web.crawled = false;
links.push(web);
if let Some(mut parsed) = url {
parsed.set_query(None);
parsed.set_fragment(None);
debug!("Final cleaned URL: `{}`", parsed.to_string());
let web = Website::new(&parsed.to_string(), false);
links.push(web);
}
}
}
return TokenSinkResult::Script(links);
}
local_name!("button") | local_name!("meta") | local_name!("iframe") => {
@@ -69,22 +66,21 @@ impl TokenSink for Website {
pub async fn parse(site: &Website, data: &str) -> Vec<Website> {
// prep work
let mut other_sites: Vec<Website> = Vec::new();
let _t = Timer::start("Parsed page");
// change data into something that can be tokenized
let chunk = Tendril::from_str(&data).expect("Failed to parse string into Tendril!");
let chunk = Tendril::from_str(data).expect("Failed to parse string into Tendril!");
// create buffer of tokens and push our input into it
let mut token_buffer = BufferQueue::default();
let token_buffer = BufferQueue::default();
token_buffer.push_back(
chunk
.try_reinterpret::<fmt::UTF8>()
.expect("Failed to reinterprt chunk!"),
.expect("Failed to reinterpret chunk!"),
);
// create the tokenizer
let tokenizer = Tokenizer::new(site.clone(), TokenizerOpts::default());
// go thru buffer
while let TokenizerResult::Script(mut sites) = tokenizer.feed(&mut token_buffer) {
while let TokenizerResult::Script(mut sites) = tokenizer.feed(&token_buffer) {
other_sites.append(&mut sites);
// other_sites.push(sites);
}
@@ -93,3 +89,57 @@ pub async fn parse(site: &Website, data: &str) -> Vec<Website> {
other_sites
}
#[instrument]
fn try_get_url(parent: &Url, link: &str) -> Option<Url> {
match Url::parse(link) {
Ok(ok) => Some(ok),
Err(e) => {
if link.starts_with('#') {
trace!("Rejecting # url");
None
} else if link.starts_with("//") {
// if a url starts with "//" is assumed that it will adopt
// the same scheme as it's parent
// https://stackoverflow.com/questions/9646407/two-forward-slashes-in-a-url-src-href-attribute
let scheme = parent.scheme();
match Url::parse(&format!("{scheme}://{}", link)) {
Ok(url) => Some(url),
Err(err) => {
error!("Failed parsing realative scheme url: {}", err);
None
}
}
} else {
// # This is some sort of realative url, gonna try patching it up into an absolute
// url
match e {
url::ParseError::RelativeUrlWithoutBase => {
// Is: scheme://host:port
let origin = parent.origin().ascii_serialization();
let url = origin.clone() + link;
trace!("Built `{url}` from `{origin} + {}`", link.to_string());
if let Ok(url) = Url::parse(&url) {
trace!("Saved relative url `{}` AS: `{}`", link, url);
Some(url)
} else {
error!(
"Failed to reconstruct a url from relative url: `{}` on site: `{}`",
link,
parent.to_string()
);
None
}
}
_ => {
error!("MISC error: {:?} {:?}", e, link);
None
}
}
}
}
}
}

101
src/s3.rs
View File

@@ -1,101 +0,0 @@
use metrics::counter;
use minio::s3::{
args::{BucketExistsArgs, MakeBucketArgs},
client::ClientBuilder,
creds::StaticProvider,
error::Error,
http::BaseUrl,
Client,
};
use tracing::{instrument, trace, warn};
use url::Url;
use crate::{db::Website, Config, Timer};
const S3_ROUND_TRIP_METRIC: &'static str = "s3_trips";
#[derive(Clone)]
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>()
.expect("Failed to parse url into BaseUrl");
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)
.expect("Failed to check if bucket exists"),
)
.await?;
counter!(S3_ROUND_TRIP_METRIC).increment(1);
if !exists {
trace!("Creating bucket...");
client
.make_bucket(
&MakeBucketArgs::new(&config.s3_bucket).expect("Failed to create bucket!"),
)
.await?;
}
counter!(S3_ROUND_TRIP_METRIC).increment(1);
trace!("Connection successful");
Ok(Self {
bucket_name: config.s3_bucket.to_owned(),
client,
})
}
#[instrument(name = "s3_store", skip_all)]
pub async fn store(&self, data: &str, url: &Url) {
let counter = counter!(S3_ROUND_TRIP_METRIC);
let t = Timer::start("Stored page");
let _ = t; // prevent compiler drop
let filename = Website::get_url_as_string(url);
trace!("Storing {} as {filename}", url.to_string());
counter.increment(1);
let _ = match &self
.client
.put_object_content(&self.bucket_name, &filename, data.to_owned())
.send()
.await
{
Ok(_) => {}
Err(err) => match err {
Error::InvalidObjectName(_) => {
// This code will really only run if the url has non-english chars
warn!("Tried storing invalid object name, retrying with Base64 encoding. Last try.");
let filename: String = Website::get_url_as_b64_path(url);
counter.increment(1);
let _ = &self
.client
.put_object_content(&self.bucket_name, &filename, data.to_owned())
.send()
.await
.unwrap();
}
_ => {}
},
};
}
}