lookup.rs

use std::collections::HashMap;

use itertools::Itertools;
use sqlx::{query, Pool, Postgres};

#[derive(Debug)]
pub struct DomainView {
    pub query: String,
    pub links_from: Vec<(String, String)>,
    pub links_to: Vec<(String, String)>,
}

impl DomainView {
    pub async fn lookup(pool: &Pool<Postgres>, domain: &str) -> anyhow::Result<Self> {
        let domain = domain.to_lowercase();

        let wildcard = domain.starts_with('*');

        let mut query: String = domain
            .chars()
            .filter(|x| x.is_ascii_alphanumeric() || *x == '.')
            .collect();

        // yes I realize it's dumb we take this off and put it back after
        // but we need this gone for the SQL query to do as we expect
        if query.starts_with('.') {
            query.remove(0);
        }

        let reversed = reverse_domain(&query);

        if wildcard {
            if !query.starts_with('.') {
                query.insert(0, '.');
            }

            query.insert(0, '*');
        }

        let nodes = query!(
            r#"
SELECT id
FROM nodes
WHERE name = $1 OR ($2 AND name LIKE $1 || '.%')
LIMIT 1000
"#,
            reversed,
            wildcard
        )
        .fetch_all(pool)
        .await?;

        let node_ids: Vec<i32> = nodes.into_iter().map(|x| x.id).collect();

        let edges_to = query!(
            "SELECT from_id, to_id FROM edges WHERE to_id = ANY($1::INT[]) LIMIT 5000",
            &node_ids
        )
        .fetch_all(pool)
        .await?;

        let edges_from = query!(
            "SELECT from_id, to_id FROM edges WHERE from_id = ANY($1::INT[]) LIMIT 5000",
            &node_ids
        )
        .fetch_all(pool)
        .await?;

        let mut ids: Vec<_> = edges_to
            .iter()
            .flat_map(|x| [x.from_id, x.to_id])
            .chain(edges_from.iter().flat_map(|x| [x.from_id, x.to_id]))
            .collect();
        ids.dedup();

        let all_nodes: HashMap<_, _> =
            query!("SELECT id, name FROM nodes WHERE id = ANY($1::INT[])", &ids)
                .fetch_all(pool)
                .await?
                .into_iter()
                .map(|x| (x.id, x.name))
                .collect();

        Ok(Self {
            query,
            links_from: edges_from
                .into_iter()
                .map(|rec| {
                    (
                        reverse_domain(&all_nodes[&rec.from_id]),
                        reverse_domain(&all_nodes[&rec.to_id]),
                    )
                })
                .collect(),
            links_to: edges_to
                .into_iter()
                .map(|rec| {
                    (
                        reverse_domain(&all_nodes[&rec.from_id]),
                        reverse_domain(&all_nodes[&rec.to_id]),
                    )
                })
                .collect(),
        })
    }

    pub fn is_empty(&self) -> bool {
        self.links_from.is_empty() && self.links_to.is_empty()
    }
}

fn reverse_domain(d: &str) -> String {
    d.split('.').rev().join(".")
}