126 lines
3.3 KiB
Rust
126 lines
3.3 KiB
Rust
use metrics::counter;
|
|
use std::fmt::Debug;
|
|
use serde::{Deserialize, Serialize};
|
|
use surrealdb::{
|
|
engine::remote::ws::{Client, Ws},
|
|
opt::auth::Root,
|
|
sql::Thing,
|
|
Surreal,
|
|
};
|
|
use tracing::{error, instrument, trace};
|
|
use url::Url;
|
|
|
|
use crate::Config;
|
|
|
|
const STORE: &str = "surql_store_calls";
|
|
|
|
#[derive(Serialize, Deserialize, Clone, Eq, PartialEq, Hash)]
|
|
pub struct Website {
|
|
/// The url that this data is found at
|
|
pub site: Url,
|
|
/// Wether or not this link has been crawled yet
|
|
pub crawled: bool,
|
|
/// 200, 404, etc
|
|
pub status_code: u16,
|
|
}
|
|
|
|
// manual impl to make tracing look nicer
|
|
impl Debug for Website {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
f.debug_struct("Website")
|
|
.field("host", &self.site.host())
|
|
.field("path", &self.site.path())
|
|
.field("status_code", &self.status_code)
|
|
.finish()
|
|
}
|
|
}
|
|
|
|
impl Website {
|
|
/// Creates a blank site (assumes that url param is site's root)
|
|
pub fn new(url: &str, crawled: bool) -> Self {
|
|
let site = match Url::parse(url) {
|
|
Ok(a) => a,
|
|
Err(_) => todo!(),
|
|
};
|
|
Self {
|
|
crawled,
|
|
site,
|
|
status_code: 0,
|
|
}
|
|
}
|
|
|
|
// 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.
|
|
#[instrument(skip(db))]
|
|
pub async fn store_all(all: Vec<Self>, db: &Surreal<Client>) -> Vec<Thing> {
|
|
counter!(STORE).increment(1);
|
|
let mut things = Vec::with_capacity(all.len());
|
|
|
|
// FIXME failes *sometimes* because "Resource Busy"
|
|
match db
|
|
.query(
|
|
"INSERT INTO website $array
|
|
ON DUPLICATE KEY UPDATE
|
|
accessed_at = time::now(),
|
|
status_code = $input.status_code,
|
|
crawled = crawled OR $input.crawled
|
|
RETURN VALUE id;
|
|
",
|
|
)
|
|
.bind(("array", all))
|
|
.await
|
|
{
|
|
Ok(mut id) => match id.take::<Vec<Thing>>(0) {
|
|
Ok(mut x) => things.append(&mut x),
|
|
Err(err) => error!("{:?}", err),
|
|
},
|
|
Err(err) => {
|
|
error!("{:?}", err);
|
|
}
|
|
}
|
|
things
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
pub struct Email {
|
|
pub email: String,
|
|
pub on: String,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
pub struct Record {
|
|
#[allow(dead_code)]
|
|
pub id: Thing,
|
|
}
|
|
|
|
#[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>(&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)
|
|
}
|
|
|