mirror of
https://github.com/Syfaro/fuzzysearch.git
synced 2024-11-10 17:02:38 +00:00
Twitter, add id_str field.
This commit is contained in:
parent
cfd05672dd
commit
cad0016522
@ -1,6 +1,6 @@
|
|||||||
use crate::models::image_query;
|
use crate::models::image_query;
|
||||||
use crate::types::*;
|
use crate::types::*;
|
||||||
use crate::utils::{extract_e621_rows, extract_fa_rows};
|
use crate::utils::{extract_e621_rows, extract_fa_rows, extract_twitter_rows};
|
||||||
use crate::{rate_limit, Pool};
|
use crate::{rate_limit, Pool};
|
||||||
use log::{debug, info};
|
use log::{debug, info};
|
||||||
use warp::{reject, Rejection, Reply};
|
use warp::{reject, Rejection, Reply};
|
||||||
@ -79,25 +79,27 @@ pub async fn search_image(
|
|||||||
|
|
||||||
debug!("Matching hash {}", num);
|
debug!("Matching hash {}", num);
|
||||||
|
|
||||||
let (fa_results, e621_results) = {
|
let results = {
|
||||||
if opts.search_type == Some(ImageSearchType::Force) {
|
if opts.search_type == Some(ImageSearchType::Force) {
|
||||||
image_query(&db, vec![num], 10).await.unwrap()
|
image_query(&db, vec![num], 10).await.unwrap()
|
||||||
} else {
|
} else {
|
||||||
let (fa_results, e621_results) = image_query(&db, vec![num], 0).await.unwrap();
|
let results = image_query(&db, vec![num], 0).await.unwrap();
|
||||||
if fa_results.len() + e621_results.len() == 0
|
if results.is_empty() && opts.search_type != Some(ImageSearchType::Exact) {
|
||||||
&& opts.search_type != Some(ImageSearchType::Exact)
|
|
||||||
{
|
|
||||||
image_query(&db, vec![num], 10).await.unwrap()
|
image_query(&db, vec![num], 10).await.unwrap()
|
||||||
} else {
|
} else {
|
||||||
(fa_results, e621_results)
|
results
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
let mut items = Vec::with_capacity(fa_results.len() + e621_results.len());
|
let mut items = Vec::with_capacity(results.len());
|
||||||
|
|
||||||
items.extend(extract_fa_rows(fa_results, Some(&hash.as_bytes())));
|
items.extend(extract_fa_rows(results.furaffinity, Some(&hash.as_bytes())));
|
||||||
items.extend(extract_e621_rows(e621_results, Some(&hash.as_bytes())));
|
items.extend(extract_e621_rows(results.e621, Some(&hash.as_bytes())));
|
||||||
|
items.extend(extract_twitter_rows(
|
||||||
|
results.twitter,
|
||||||
|
Some(&hash.as_bytes()),
|
||||||
|
));
|
||||||
|
|
||||||
items.sort_by(|a, b| {
|
items.sort_by(|a, b| {
|
||||||
a.distance
|
a.distance
|
||||||
@ -124,6 +126,7 @@ pub async fn search_hashes(
|
|||||||
let hashes: Vec<i64> = opts
|
let hashes: Vec<i64> = opts
|
||||||
.hashes
|
.hashes
|
||||||
.split(',')
|
.split(',')
|
||||||
|
.take(10)
|
||||||
.filter_map(|hash| hash.parse::<i64>().ok())
|
.filter_map(|hash| hash.parse::<i64>().ok())
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
@ -133,13 +136,14 @@ pub async fn search_hashes(
|
|||||||
|
|
||||||
rate_limit!(&api_key, &db, image_limit, "image", hashes.len() as i16);
|
rate_limit!(&api_key, &db, image_limit, "image", hashes.len() as i16);
|
||||||
|
|
||||||
let (fa_matches, e621_matches) = image_query(&db, hashes, 10)
|
let results = image_query(&db, hashes, 10)
|
||||||
.await
|
.await
|
||||||
.map_err(|err| reject::custom(Error::from(err)))?;
|
.map_err(|err| reject::custom(Error::from(err)))?;
|
||||||
|
|
||||||
let mut matches = Vec::with_capacity(fa_matches.len() + e621_matches.len());
|
let mut matches = Vec::with_capacity(results.len());
|
||||||
matches.extend(extract_fa_rows(fa_matches, None));
|
matches.extend(extract_fa_rows(results.furaffinity, None));
|
||||||
matches.extend(extract_e621_rows(e621_matches, None));
|
matches.extend(extract_e621_rows(results.e621, None));
|
||||||
|
matches.extend(extract_twitter_rows(results.twitter, None));
|
||||||
|
|
||||||
Ok(warp::reply::json(&matches))
|
Ok(warp::reply::json(&matches))
|
||||||
}
|
}
|
||||||
@ -190,6 +194,7 @@ pub async fn search_file(
|
|||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|row| File {
|
.map(|row| File {
|
||||||
id: row.get("id"),
|
id: row.get("id"),
|
||||||
|
id_str: row.get::<&str, i32>("id").to_string(),
|
||||||
url: row.get("url"),
|
url: row.get("url"),
|
||||||
filename: row.get("filename"),
|
filename: row.get("filename"),
|
||||||
artists: row
|
artists: row
|
||||||
|
@ -35,24 +35,43 @@ pub async fn lookup_api_key(key: &str, db: DB<'_>) -> Option<ApiKey> {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub struct ImageQueryResults {
|
||||||
|
pub furaffinity: Vec<tokio_postgres::Row>,
|
||||||
|
pub e621: Vec<tokio_postgres::Row>,
|
||||||
|
pub twitter: Vec<tokio_postgres::Row>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ImageQueryResults {
|
||||||
|
#[inline]
|
||||||
|
pub fn len(&self) -> usize {
|
||||||
|
self.furaffinity.len() + self.e621.len() + self.twitter.len()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[inline]
|
||||||
|
pub fn is_empty(&self) -> bool {
|
||||||
|
self.len() == 0
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub async fn image_query(
|
pub async fn image_query(
|
||||||
db: DB<'_>,
|
db: DB<'_>,
|
||||||
hashes: Vec<i64>,
|
hashes: Vec<i64>,
|
||||||
distance: i64,
|
distance: i64,
|
||||||
) -> Result<(Vec<tokio_postgres::Row>, Vec<tokio_postgres::Row>), tokio_postgres::Error> {
|
) -> Result<ImageQueryResults, tokio_postgres::Error> {
|
||||||
let mut params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> =
|
let mut params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> =
|
||||||
Vec::with_capacity(hashes.len() + 1);
|
Vec::with_capacity(hashes.len() + 1);
|
||||||
params.insert(0, &distance);
|
params.insert(0, &distance);
|
||||||
|
|
||||||
let mut fa_where_clause = Vec::with_capacity(hashes.len());
|
let mut fa_where_clause = Vec::with_capacity(hashes.len());
|
||||||
let mut e621_where_clause = Vec::with_capacity(hashes.len());
|
let mut hash_where_clause = Vec::with_capacity(hashes.len());
|
||||||
|
|
||||||
for (idx, hash) in hashes.iter().enumerate() {
|
for (idx, hash) in hashes.iter().enumerate() {
|
||||||
params.push(hash);
|
params.push(hash);
|
||||||
|
|
||||||
fa_where_clause.push(format!(" hash_int <@ (${}, $1)", idx + 2));
|
fa_where_clause.push(format!(" hash_int <@ (${}, $1)", idx + 2));
|
||||||
e621_where_clause.push(format!(" hash <@ (${}, $1)", idx + 2));
|
hash_where_clause.push(format!(" hash <@ (${}, $1)", idx + 2));
|
||||||
}
|
}
|
||||||
|
let hash_where_clause = hash_where_clause.join(" OR ");
|
||||||
|
|
||||||
let fa_query = format!(
|
let fa_query = format!(
|
||||||
"SELECT
|
"SELECT
|
||||||
@ -93,12 +112,30 @@ pub async fn image_query(
|
|||||||
) artists
|
) artists
|
||||||
WHERE
|
WHERE
|
||||||
{}",
|
{}",
|
||||||
e621_where_clause.join(" OR ")
|
&hash_where_clause
|
||||||
);
|
);
|
||||||
|
|
||||||
let fa = db.query::<str>(&*fa_query, ¶ms);
|
let twitter_query = format!(
|
||||||
let e621 = db.query::<str>(&*e621_query, ¶ms);
|
"SELECT
|
||||||
|
twitter_view.id,
|
||||||
|
twitter_view.artists,
|
||||||
|
twitter_view.url,
|
||||||
|
twitter_view.hash
|
||||||
|
FROM
|
||||||
|
twitter_view
|
||||||
|
WHERE
|
||||||
|
{}",
|
||||||
|
&hash_where_clause
|
||||||
|
);
|
||||||
|
|
||||||
let results = futures::future::join(fa, e621).await;
|
let furaffinity = db.query::<str>(&*fa_query, ¶ms);
|
||||||
Ok((results.0?, results.1?))
|
let e621 = db.query::<str>(&*e621_query, ¶ms);
|
||||||
|
let twitter = db.query::<str>(&*twitter_query, ¶ms);
|
||||||
|
|
||||||
|
let results = futures::future::join3(furaffinity, e621, twitter).await;
|
||||||
|
Ok(ImageQueryResults {
|
||||||
|
furaffinity: results.0?,
|
||||||
|
e621: results.1?,
|
||||||
|
twitter: results.2?,
|
||||||
|
})
|
||||||
}
|
}
|
||||||
|
@ -25,7 +25,8 @@ pub enum RateLimit {
|
|||||||
/// A general type for every file.
|
/// A general type for every file.
|
||||||
#[derive(Debug, Default, Serialize)]
|
#[derive(Debug, Default, Serialize)]
|
||||||
pub struct File {
|
pub struct File {
|
||||||
pub id: i32,
|
pub id: i64,
|
||||||
|
pub id_str: String,
|
||||||
pub url: String,
|
pub url: String,
|
||||||
pub filename: String,
|
pub filename: String,
|
||||||
pub artists: Option<Vec<String>>,
|
pub artists: Option<Vec<String>>,
|
||||||
@ -46,6 +47,7 @@ pub enum SiteInfo {
|
|||||||
FurAffinity(FurAffinityFile),
|
FurAffinity(FurAffinityFile),
|
||||||
#[serde(rename = "e621")]
|
#[serde(rename = "e621")]
|
||||||
E621(E621File),
|
E621(E621File),
|
||||||
|
Twitter,
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Information about a file hosted on FurAffinity.
|
/// Information about a file hosted on FurAffinity.
|
||||||
|
40
src/utils.rs
40
src/utils.rs
@ -76,7 +76,8 @@ pub fn extract_fa_rows<'a>(
|
|||||||
let dbbytes: Vec<u8> = row.get("hash");
|
let dbbytes: Vec<u8> = row.get("hash");
|
||||||
|
|
||||||
File {
|
File {
|
||||||
id: row.get("id"),
|
id: row.get::<&str, i32>("id") as i64,
|
||||||
|
id_str: row.get::<&str, i32>("id").to_string(),
|
||||||
url: row.get("url"),
|
url: row.get("url"),
|
||||||
filename: row.get("filename"),
|
filename: row.get("filename"),
|
||||||
hash: row.get("hash_int"),
|
hash: row.get("hash_int"),
|
||||||
@ -100,7 +101,8 @@ pub fn extract_e621_rows<'a>(
|
|||||||
let dbbytes = dbhash.to_be_bytes();
|
let dbbytes = dbhash.to_be_bytes();
|
||||||
|
|
||||||
File {
|
File {
|
||||||
id: row.get("id"),
|
id: row.get::<&str, i32>("id") as i64,
|
||||||
|
id_str: row.get::<&str, i32>("id").to_string(),
|
||||||
url: row.get("url"),
|
url: row.get("url"),
|
||||||
hash: Some(dbhash),
|
hash: Some(dbhash),
|
||||||
distance: hash
|
distance: hash
|
||||||
@ -115,3 +117,37 @@ pub fn extract_e621_rows<'a>(
|
|||||||
}
|
}
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn extract_twitter_rows<'a>(
|
||||||
|
rows: Vec<tokio_postgres::Row>,
|
||||||
|
hash: Option<&'a [u8]>,
|
||||||
|
) -> impl IntoIterator<Item = File> + 'a {
|
||||||
|
rows.into_iter().map(move |row| {
|
||||||
|
let dbhash: i64 = row.get("hash");
|
||||||
|
let dbbytes = dbhash.to_be_bytes();
|
||||||
|
|
||||||
|
let url: String = row.get("url");
|
||||||
|
|
||||||
|
let filename = url
|
||||||
|
.split('/')
|
||||||
|
.last()
|
||||||
|
.unwrap()
|
||||||
|
.split(':')
|
||||||
|
.next()
|
||||||
|
.unwrap()
|
||||||
|
.to_string();
|
||||||
|
|
||||||
|
File {
|
||||||
|
id: row.get("id"),
|
||||||
|
id_str: row.get::<&str, i64>("id").to_string(),
|
||||||
|
url,
|
||||||
|
hash: Some(dbhash),
|
||||||
|
distance: hash
|
||||||
|
.map(|hash| hamming::distance_fast(&dbbytes, &hash).ok())
|
||||||
|
.flatten(),
|
||||||
|
site_info: Some(SiteInfo::Twitter),
|
||||||
|
artists: row.get("artists"),
|
||||||
|
filename,
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
Loading…
Reference in New Issue
Block a user