mirror of
https://github.com/Syfaro/fuzzysearch.git
synced 2024-11-10 17:02:38 +00:00
Some improvements.
This commit is contained in:
parent
70b94f454b
commit
658a023334
@ -20,6 +20,7 @@ pub fn search_image(db: Pool) -> impl Filter<Extract = impl Reply, Error = Rejec
|
|||||||
warp::path("image")
|
warp::path("image")
|
||||||
.and(warp::post())
|
.and(warp::post())
|
||||||
.and(warp::multipart::form().max_length(1024 * 1024 * 10))
|
.and(warp::multipart::form().max_length(1024 * 1024 * 10))
|
||||||
|
.and(warp::query::<ImageSearchOpts>())
|
||||||
.and(with_pool(db))
|
.and(with_pool(db))
|
||||||
.and(with_api_key())
|
.and(with_api_key())
|
||||||
.and_then(handlers::search_image)
|
.and_then(handlers::search_image)
|
||||||
|
@ -1,7 +1,8 @@
|
|||||||
|
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};
|
||||||
use crate::{rate_limit, Pool};
|
use crate::{rate_limit, Pool};
|
||||||
use log::{info, debug};
|
use log::{debug, info};
|
||||||
use warp::{reject, Rejection, Reply};
|
use warp::{reject, Rejection, Reply};
|
||||||
|
|
||||||
fn map_bb8_err(err: bb8::RunError<tokio_postgres::Error>) -> Rejection {
|
fn map_bb8_err(err: bb8::RunError<tokio_postgres::Error>) -> Rejection {
|
||||||
@ -37,6 +38,7 @@ impl warp::reject::Reject for Error {}
|
|||||||
|
|
||||||
pub async fn search_image(
|
pub async fn search_image(
|
||||||
form: warp::multipart::FormData,
|
form: warp::multipart::FormData,
|
||||||
|
opts: ImageSearchOpts,
|
||||||
db: Pool,
|
db: Pool,
|
||||||
api_key: String,
|
api_key: String,
|
||||||
) -> Result<impl Reply, Rejection> {
|
) -> Result<impl Reply, Rejection> {
|
||||||
@ -77,52 +79,20 @@ pub async fn search_image(
|
|||||||
|
|
||||||
debug!("Matching hash {}", num);
|
debug!("Matching hash {}", num);
|
||||||
|
|
||||||
let params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = vec![&num];
|
let (fa_results, e621_results) = {
|
||||||
|
if opts.search_type == Some(ImageSearchType::Force) {
|
||||||
let fa = db.query(
|
image_query(&db, num, 10).await.unwrap()
|
||||||
"SELECT
|
} else {
|
||||||
submission.id,
|
let (fa_results, e621_results) = image_query(&db, num, 0).await.unwrap();
|
||||||
submission.url,
|
if fa_results.len() + e621_results.len() == 0
|
||||||
submission.filename,
|
&& opts.search_type != Some(ImageSearchType::Exact)
|
||||||
submission.file_id,
|
{
|
||||||
submission.hash,
|
image_query(&db, num, 10).await.unwrap()
|
||||||
submission.hash_int,
|
} else {
|
||||||
artist.name
|
(fa_results, e621_results)
|
||||||
FROM
|
}
|
||||||
submission
|
}
|
||||||
JOIN artist
|
};
|
||||||
ON artist.id = submission.artist_id
|
|
||||||
WHERE
|
|
||||||
hash_int <@ ($1, 10)",
|
|
||||||
¶ms,
|
|
||||||
);
|
|
||||||
|
|
||||||
let e621 = db.query(
|
|
||||||
"SELECT
|
|
||||||
e621.id,
|
|
||||||
e621.hash,
|
|
||||||
e621.data->>'file_url' url,
|
|
||||||
e621.data->>'md5' md5,
|
|
||||||
sources.list sources,
|
|
||||||
artists.list artists,
|
|
||||||
(e621.data->>'md5') || '.' || (e621.data->>'file_ext') filename
|
|
||||||
FROM
|
|
||||||
e621,
|
|
||||||
LATERAL (
|
|
||||||
SELECT array_agg(s) list
|
|
||||||
FROM jsonb_array_elements_text(data->'sources') s
|
|
||||||
) sources,
|
|
||||||
LATERAL (
|
|
||||||
SELECT array_agg(s) list
|
|
||||||
FROM jsonb_array_elements_text(data->'artist') s
|
|
||||||
) artists
|
|
||||||
WHERE
|
|
||||||
hash <@ ($1, 10)",
|
|
||||||
¶ms,
|
|
||||||
);
|
|
||||||
|
|
||||||
let results = futures::future::join(fa, e621).await;
|
|
||||||
let (fa_results, e621_results) = (results.0.unwrap(), results.1.unwrap());
|
|
||||||
|
|
||||||
let mut items = Vec::with_capacity(fa_results.len() + e621_results.len());
|
let mut items = Vec::with_capacity(fa_results.len() + e621_results.len());
|
||||||
|
|
||||||
|
@ -27,9 +27,12 @@ async fn main() {
|
|||||||
let log = warp::log("fuzzysearch");
|
let log = warp::log("fuzzysearch");
|
||||||
let cors = warp::cors()
|
let cors = warp::cors()
|
||||||
.allow_any_origin()
|
.allow_any_origin()
|
||||||
|
.allow_headers(vec!["x-api-key"])
|
||||||
.allow_methods(vec!["GET", "POST"]);
|
.allow_methods(vec!["GET", "POST"]);
|
||||||
|
|
||||||
let api = filters::search(db_pool);
|
let options = warp::options().map(|| "✓");
|
||||||
|
|
||||||
|
let api = options.or(filters::search(db_pool));
|
||||||
let routes = api
|
let routes = api
|
||||||
.or(warp::path::end()
|
.or(warp::path::end()
|
||||||
.map(|| warp::redirect(warp::http::Uri::from_static("https://fuzzysearch.net"))))
|
.map(|| warp::redirect(warp::http::Uri::from_static("https://fuzzysearch.net"))))
|
||||||
|
@ -34,3 +34,56 @@ pub async fn lookup_api_key(key: &str, db: DB<'_>) -> Option<ApiKey> {
|
|||||||
_ => None,
|
_ => None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub async fn image_query(
|
||||||
|
db: DB<'_>,
|
||||||
|
num: i64,
|
||||||
|
distance: i64,
|
||||||
|
) -> Result<(Vec<tokio_postgres::Row>, Vec<tokio_postgres::Row>), tokio_postgres::Error> {
|
||||||
|
let params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = vec![&num, &distance];
|
||||||
|
|
||||||
|
let fa = db.query(
|
||||||
|
"SELECT
|
||||||
|
submission.id,
|
||||||
|
submission.url,
|
||||||
|
submission.filename,
|
||||||
|
submission.file_id,
|
||||||
|
submission.hash,
|
||||||
|
submission.hash_int,
|
||||||
|
artist.name
|
||||||
|
FROM
|
||||||
|
submission
|
||||||
|
JOIN artist
|
||||||
|
ON artist.id = submission.artist_id
|
||||||
|
WHERE
|
||||||
|
hash_int <@ ($1, $2)",
|
||||||
|
¶ms,
|
||||||
|
);
|
||||||
|
|
||||||
|
let e621 = db.query(
|
||||||
|
"SELECT
|
||||||
|
e621.id,
|
||||||
|
e621.hash,
|
||||||
|
e621.data->>'file_url' url,
|
||||||
|
e621.data->>'md5' md5,
|
||||||
|
sources.list sources,
|
||||||
|
artists.list artists,
|
||||||
|
(e621.data->>'md5') || '.' || (e621.data->>'file_ext') filename
|
||||||
|
FROM
|
||||||
|
e621,
|
||||||
|
LATERAL (
|
||||||
|
SELECT array_agg(s) list
|
||||||
|
FROM jsonb_array_elements_text(data->'sources') s
|
||||||
|
) sources,
|
||||||
|
LATERAL (
|
||||||
|
SELECT array_agg(s) list
|
||||||
|
FROM jsonb_array_elements_text(data->'artist') s
|
||||||
|
) artists
|
||||||
|
WHERE
|
||||||
|
hash <@ ($1, $2)",
|
||||||
|
¶ms,
|
||||||
|
);
|
||||||
|
|
||||||
|
let results = futures::future::join(fa, e621).await;
|
||||||
|
Ok((results.0?, results.1?))
|
||||||
|
}
|
||||||
|
14
src/types.rs
14
src/types.rs
@ -68,6 +68,20 @@ pub struct FileSearchOpts {
|
|||||||
pub url: Option<String>,
|
pub url: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
pub struct ImageSearchOpts {
|
||||||
|
#[serde(rename = "type")]
|
||||||
|
pub search_type: Option<ImageSearchType>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize, PartialEq)]
|
||||||
|
#[serde(rename_all = "lowercase")]
|
||||||
|
pub enum ImageSearchType {
|
||||||
|
Close,
|
||||||
|
Exact,
|
||||||
|
Force,
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Serialize)]
|
#[derive(Debug, Serialize)]
|
||||||
pub struct ImageSimilarity {
|
pub struct ImageSimilarity {
|
||||||
pub hash: i64,
|
pub hash: i64,
|
||||||
|
Loading…
Reference in New Issue
Block a user