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

formatting and timer changes, consolidated functions

Browse Source
This commit is contained in:
oliver 2024-11-12 18:40:10 -07:00
parent 8a5ac61b26
commit d28d18de08
1 changed files with 82 additions and 62 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

144
src/main.rs
View File

@ -1,12 +1,14 @@
extern crate markup5ever_rcdom as rcdom;
extern crate html5ever; extern crate html5ever;
extern crate markup5ever_rcdom as rcdom;
use std::{path::is_separator, rc::Rc, time::Instant};
use db::{connect, Website}; use db::{connect, Website};
use html5ever::{local_name, parse_document, tendril::TendrilSink, tree_builder::TreeBuilderOpts, ParseOpts}; use html5ever::{
use rcdom::{Node, RcDom}; local_name, parse_document, tendril::TendrilSink, tree_builder::TreeBuilderOpts, ParseOpts,
};
use rcdom::RcDom;
use std::time::Instant;
use surrealdb::{engine::remote::ws::Client, sql::Thing, Surreal}; use surrealdb::{engine::remote::ws::Client, sql::Thing, Surreal};
use tracing::{debug, info, instrument, trace, trace_span, warn}; use tracing::{debug, info, instrument, trace, trace_span};
use tracing_subscriber::EnvFilter; use tracing_subscriber::EnvFilter;
mod db; mod db;
@ -21,9 +23,9 @@ async fn main() {
debug!("Starting..."); debug!("Starting...");
// Would probably take these in as parameters from a cli // Would probably take these in as parameters from a cli
// let url = "https://oliveratkinson.net/"; let url = "https://oliveratkinson.net/";
let url = "http://localhost:5500"; // let url = "http://localhost:5500";
let budget = 50; let budget = 1000;
let mut crawled = 0; let mut crawled = 0;
let db = connect().await.expect("Failed to connect to db, aborting."); let db = connect().await.expect("Failed to connect to db, aborting.");
@ -39,9 +41,7 @@ async fn main() {
let pre_loop_span = span.enter(); let pre_loop_span = span.enter();
// Download the site // Download the site
let mut site = Website::new(&url, false); let mut site = Website::new(&url, false);
let dom = get(&mut site, &db, &client).await.expect("Inital page returned None."); get(&mut site, &db, &client, &mut crawled).await;
crawl_wrapper(&dom, &db, &site, &mut crawled).await;
drop(pre_loop_span); drop(pre_loop_span);
@ -50,7 +50,9 @@ async fn main() {
while crawled < budget { while crawled < budget {
let get_num = if budget - crawled < 100 { let get_num = if budget - crawled < 100 {
budget - crawled budget - crawled
} else {100}; } else {
100
};
let uncrawled = get_uncrawled_links(&db, get_num).await; let uncrawled = get_uncrawled_links(&db, get_num).await;
if uncrawled.len() == 0 { if uncrawled.len() == 0 {
@ -63,13 +65,9 @@ async fn main() {
let _ = span.enter(); let _ = span.enter();
for mut site in uncrawled { for mut site in uncrawled {
if let Some(dom) = get(&mut site, &db, &client).await { get(&mut site, &db, &client, &mut crawled).await;
crawl_wrapper(&dom, &db, &site, &mut crawled).await; let percent = format!("{:.2}%", (crawled as f32 / budget as f32) * 100f32);
let percent = format!("{:.2}%", (crawled as f32/budget as f32) * 100f32); info!("Crawled {crawled} out of {budget} pages. ({percent})");
info!("Crawled {crawled} out of {budget} pages. ({percent})");
} else {
warn!("Failed to get {}", site.to_string());
}
} }
} }
drop(span); drop(span);
@ -77,24 +75,21 @@ async fn main() {
info!("Done"); info!("Done");
} }
async fn crawl_wrapper(dom: &Rc<Node>, db: &Surreal<Client>, site: &Website, count: &mut usize) {
let mut buffer = Vec::new();
let timer= Timer::start("Walked");
walk(&dom, &db, &site, &mut buffer).await;
drop(timer);
site.links_to(buffer, &db).await;
*count += 1;
}
#[instrument(skip_all)] #[instrument(skip_all)]
/// A quick helper function for downloading a url /// A quick helper function for downloading a url
async fn get(site: &mut Website, db: &Surreal<Client>, getter: &reqwest::Client) -> Option<Rc<Node>> { async fn get(
site: &mut Website,
db: &Surreal<Client>,
request_client: &reqwest::Client,
count: &mut usize,
) {
trace!("Get: {}", site.to_string()); trace!("Get: {}", site.to_string());
let timer = Timer::start("Got page"); let timer = Timer::start("Got page");
if let Ok(response) = getter.get(site.to_string()).send().await { if let Ok(response) = request_client.get(site.to_string()).send().await {
drop(timer); timer.stop();
// Get body
let data = response.text().await.unwrap(); let data = response.text().await.unwrap();
let opts = ParseOpts { let opts = ParseOpts {
tree_builder: TreeBuilderOpts { tree_builder: TreeBuilderOpts {
@ -103,46 +98,63 @@ async fn get(site: &mut Website, db: &Surreal<Client>, getter: &reqwest::Client)
}, },
..Default::default() ..Default::default()
}; };
// Get DOM
let dom = parse_document(RcDom::default(), opts) let dom = parse_document(RcDom::default(), opts)
.from_utf8() .from_utf8()
.read_from(&mut data.as_bytes()) .read_from(&mut data.as_bytes())
.unwrap(); .unwrap();
// TODO save the dom to minio if a flag is set // TODO save the dom to minio if a flag is set
// Modify record in database
site.set_crawled(); site.set_crawled();
site.store(db).await; site.store(db).await;
trace!("Got: {}", site.to_string()); trace!("Got: {}", site.to_string());
return Some(dom.document);
// Walk all the children nodes, searching for links to other pages.
let mut buffer = Vec::new();
let timer = Timer::start("Walked");
walk(&dom.document, &db, &site, &mut buffer).await;
timer.stop();
// Put all the found links into the database.
site.links_to(buffer, &db).await;
*count += 1;
} }
trace!("Failed to get: {}", site.to_string()); trace!("Failed to get: {}", site.to_string());
None
} }
/// Walks the givin site, placing it's findings in the database /// Walks the givin site, placing it's findings in the database
async fn walk(node: &rcdom::Handle, db: &Surreal<Client> , site: &Website, links_to: &mut Vec<Thing>) { async fn walk(
node: &rcdom::Handle,
db: &Surreal<Client>,
site: &Website,
links_to: &mut Vec<Thing>,
) {
let span = trace_span!("Walk"); let span = trace_span!("Walk");
let span = span.enter(); let span = span.enter();
// Match each node - node basically means element. // Match each node - node basically means element.
match &node.data { match &node.data {
rcdom::NodeData::Element { name, attrs, template_contents, mathml_annotation_xml_integration_point } => { rcdom::NodeData::Element { name, attrs, .. } => {
for attr in attrs.borrow().clone() { for attr in attrs.borrow().clone() {
match name.local { match name.local {
local_name!("a") | local_name!("a")
local_name!("audio") | | local_name!("audio")
local_name!("area") | | local_name!("area")
local_name!("img") | | local_name!("img")
local_name!("link") | | local_name!("link")
local_name!("object") | | local_name!("object")
local_name!("source") | | local_name!("source")
local_name!("base") | | local_name!("base")
local_name!("video") => { | local_name!("video") => {
let attribute_name = attr.name.local.to_string(); let attribute_name = attr.name.local.to_string();
if attribute_name == "src" || attribute_name == "href" || attribute_name == "data" { if attribute_name == "src"
|| attribute_name == "href"
|| attribute_name == "data"
{
// Get clone of the current site object // Get clone of the current site object
let mut web = site.clone(); let mut web = site.clone();
// Set url // Set url
let url = web.mut_url(); let url = web.mut_url();
url.set_fragment(None); // removes #xyz url.set_fragment(None); // removes #xyz
@ -158,21 +170,19 @@ async fn walk(node: &rcdom::Handle, db: &Surreal<Client> , site: &Website, links
links_to.push(id); links_to.push(id);
} }
} }
}, }
local_name!("button") | local_name!("button") | local_name!("meta") | local_name!("iframe") => {
local_name!("meta") |
local_name!("iframe") => {
// dbg!(attrs); // dbg!(attrs);
} }
_ => {/**/} _ => {}
}; };
}; }
}, }
_ => {}, _ => {}
}; };
drop(span); drop(span);
for child in node.children.borrow().iter() { for child in node.children.borrow().iter() {
Box::pin(walk(child, db, site, links_to)).await; Box::pin(walk(child, db, site, links_to)).await;
} }
} }
@ -187,7 +197,9 @@ async fn get_uncrawled_links(db: &Surreal<Client>, mut count: usize) -> Vec<Webs
.bind(("count", count)) .bind(("count", count))
.await .await
.expect("Hard-coded query failed..?"); .expect("Hard-coded query failed..?");
response.take(0).expect("Returned websites couldn't be parsed") response
.take(0)
.expect("Returned websites couldn't be parsed")
} }
pub struct Timer<'a> { pub struct Timer<'a> {
@ -198,13 +210,21 @@ pub struct Timer<'a> {
impl<'a> Timer<'a> { impl<'a> Timer<'a> {
#[inline] #[inline]
pub fn start(msg: &'a str) -> Self { pub fn start(msg: &'a str) -> Self {
Self { start: Instant::now(), msg } Self {
start: Instant::now(),
msg,
}
}
pub fn stop(&self) -> f64 {
let dif = self.start.elapsed().as_micros();
let ms = dif as f64 / 1000.;
debug!("{}", format!("{} in {:.3}ms", self.msg, ms));
ms
} }
} }
impl<'a> Drop for Timer<'a> { impl Drop for Timer<'_> {
fn drop(&mut self) { fn drop(&mut self) {
let dif = self.start.elapsed().as_micros(); self.stop();
debug!("{}", format!("{} in {:.3}ms", self.msg, dif as f64/1000.));
} }
} }