Oliver/internet_mapper
1
0
Fork 0
Code Issues 1 Pull Requests Actions Packages Projects Releases Wiki Activity

helper code

Browse Source
This commit is contained in:
Rushmore75 2025-07-09 15:58:22 -06:00
parent 50606bb69e
commit 6790061e22
2 changed files with 24 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
src/db.rs
View File

@ -27,7 +27,11 @@ pub struct Website {
// 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("site", &self.site).finish()
f.debug_struct("Website")
.field("host", &self.site.host())
.field("path", &self.site.path())
.field("status_code", &self.status_code)
.finish()
}
}
@ -52,6 +56,7 @@ impl Website {
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

28
src/main.rs
View File

@ -1,5 +1,4 @@
#![feature(ip_from)]
#![feature(async_closure)]
#![warn(clippy::expect_used)]
#![deny(clippy::unwrap_used)]
@ -20,7 +19,7 @@ use metrics_exporter_prometheus::PrometheusBuilder;
use serde::Deserialize;
use surrealdb::{engine::remote::ws::Client, Surreal};
use tokio::{io::{AsyncWriteExt, BufWriter}, task::JoinSet};
use tracing::{debug, error, info, instrument, level_filters::LevelFilter, trace, trace_span, warn};
use tracing::{debug, debug_span, error, info, instrument, level_filters::LevelFilter, trace, trace_span, warn};
use tracing_subscriber::{fmt, layer::SubscriberExt, EnvFilter, Layer, Registry};
mod db;
@ -180,8 +179,6 @@ async fn process(mut site: Website, db: Surreal<Client>, reqwest: reqwest::Clien
let CT = headers.get("Content-Type");
let ct = headers.get("content-type");
if CT.is_none() && ct.is_none() {
}
let ct = match (CT,ct) {
(None, None) => {
warn!("Server did not respond with Content-Type header. Url: {} Headers: ({:?})", site.site.to_string(), headers);
@ -192,19 +189,20 @@ async fn process(mut site: Website, db: Surreal<Client>, reqwest: reqwest::Clien
(Some(a), Some(_)) => a,
};
// create filepath (handles / -> /index.html)
let path = filesystem::as_path(&site.site, ct);
// make sure that the file is good to go
if let Some(file) = filesystem::init(&path).await {
let should_parse = path.to_string_lossy().ends_with(".html");
let mut buf: Vec<u8> = Vec::new();
let mut writer = BufWriter::new(file);
// Get body from response
// stream the response onto the disk
let mut stream = response.bytes_stream();
let should_parse = path.to_string_lossy().ends_with(".html");
let mut writer = BufWriter::new(file);
let mut buf: Vec<u8> = Vec::new();
// Write file to disk
info!("Writing at: {:?}", path);
while let Some(data) = stream.next().await {
match data {
@ -212,6 +210,7 @@ async fn process(mut site: Website, db: Surreal<Client>, reqwest: reqwest::Clien
let _ = writer.write_all(&data).await;
// If we are going to parse this file later, we will save it
// into memory as well as the disk.
// We do this because the data here might be incomplete
if should_parse {
data.iter().for_each(|f| buf.push(*f));
}
@ -223,7 +222,12 @@ async fn process(mut site: Website, db: Surreal<Client>, reqwest: reqwest::Clien
}
let _ = writer.flush();
// (If needed) Parse the file
if should_parse {
let span = debug_span!("Should Parse");
let enter = span.enter();
// Parse document and get relationships
let sites = parser::parse(&site, &buf).await;
// De-duplicate this list
@ -237,6 +241,8 @@ async fn process(mut site: Website, db: Surreal<Client>, reqwest: reqwest::Clien
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 _ = Website::store_all(de_dupe_sites, &db).await;
drop(enter);
}
// METRICS
@ -246,10 +252,12 @@ async fn process(mut site: Website, db: Surreal<Client>, reqwest: reqwest::Clien
// update self in db
site.crawled = true;
site.status_code = code.as_u16();
Website::store_all(vec![site], &db).await;
Website::store_all(vec![site.clone()], &db).await;
} else {
error!("File failed to cooperate: {:?}", path);
}
trace!("Done processing: {}", &site.site);
} else {
error!("Failed to get: {}", &site.site);
}