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

work off content-type header

Browse Source
This commit is contained in:
Rushmore75 2025-04-17 09:35:57 -06:00
parent 7fab961d76
commit 647c4cd324
1 changed files with 67 additions and 52 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
src/main.rs
View File

@ -19,7 +19,7 @@ use metrics::{counter, gauge};
use metrics_exporter_prometheus::PrometheusBuilder; use metrics_exporter_prometheus::PrometheusBuilder;
use serde::Deserialize; use serde::Deserialize;
use surrealdb::{engine::remote::ws::Client, Surreal}; use surrealdb::{engine::remote::ws::Client, Surreal};
use tokio::{io::AsyncWriteExt, task::JoinSet}; use tokio::{io::{AsyncWriteExt, BufWriter}, task::JoinSet};
use tracing::{debug, error, info, instrument, level_filters::LevelFilter, trace, trace_span, warn}; use tracing::{debug, error, info, instrument, level_filters::LevelFilter, trace, trace_span, warn};
use tracing_subscriber::{fmt, layer::SubscriberExt, EnvFilter, Layer, Registry}; use tracing_subscriber::{fmt, layer::SubscriberExt, EnvFilter, Layer, Registry};
@ -175,23 +175,40 @@ async fn process(mut site: Website, db: Surreal<Client>, reqwest: reqwest::Clien
if let Ok(response) = request_builder.send().await { if let Ok(response) = request_builder.send().await {
// Get body from response // Get body from response
let headers = response.headers(); let headers = response.headers();
let ct = headers.get("Content-Type");
if let Some(ct) = ct { #[allow(non_snake_case)]
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);
return
},
(None, Some(a)) => a,
(Some(a), None) => a,
(Some(a), Some(_)) => a,
};
let path = filesystem::as_path(&site.site, ct); let path = filesystem::as_path(&site.site, ct);
// make sure that the file is good to go // make sure that the file is good to go
if let Some(mut file) = filesystem::init(&path).await { if let Some(file) = filesystem::init(&path).await {
let should_parse = path.to_string_lossy().ends_with(".html"); let should_parse = path.to_string_lossy().ends_with(".html");
let mut buf: Vec<u8> = Vec::new(); let mut buf: Vec<u8> = Vec::new();
let mut writer = BufWriter::new(file);
// stream the response onto the disk // stream the response onto the disk
let mut stream = response.bytes_stream(); let mut stream = response.bytes_stream();
info!("Writing at: {:?}", path);
while let Some(data) = stream.next().await { while let Some(data) = stream.next().await {
match data { match data {
Ok(data) => { Ok(data) => {
debug!("Writing at: {:?}", path); let _ = writer.write_all(&data).await;
let _ = file.write_all(&data).await;
// If we are going to parse this file later, we will save it // If we are going to parse this file later, we will save it
// into memory as well as the disk. // into memory as well as the disk.
if should_parse { if should_parse {
@ -203,6 +220,7 @@ async fn process(mut site: Website, db: Surreal<Client>, reqwest: reqwest::Clien
}, },
} }
} }
let _ = writer.flush();
if should_parse { if should_parse {
// Parse document and get relationships // Parse document and get relationships
@ -230,9 +248,6 @@ async fn process(mut site: Website, db: Surreal<Client>, reqwest: reqwest::Clien
} else { } else {
error!("File failed to cooperate: {:?}", path); error!("File failed to cooperate: {:?}", path);
} }
} else {
warn!("Server did not respond with Content-Type header: {}", site.site.to_string());
}
} else { } else {
error!("Failed to get: {}", &site.site); error!("Failed to get: {}", &site.site);
} }