ebay_scraper_rust/src/db.rs

use chrono::{DateTime, Utc};
use rusqlite::Connection;
use serde::Serialize;
use std::path::Path;
use tracing::info;

pub trait DBTable {
    const TABLE_NAME: &'static str;
    const TABLE_SCHEMA: &'static str;
    fn initialize(conn: &Connection) {
        let create_table = &format!(
            "CREATE TABLE IF NOT EXISTS {} (
                {}
            )",
            Self::TABLE_NAME,
            Self::TABLE_SCHEMA
        );
        info!("Creating table with following schema;");
        info!("{} ({})", Self::TABLE_NAME, Self::TABLE_SCHEMA);
        conn.execute(create_table, ()).unwrap();
    }

    fn get_all(conn: &Connection) -> rusqlite::Result<Vec<Self>>
    where
        Self: Sized;
}

#[derive(Serialize, Debug, PartialEq, Clone)]
pub struct SearchURL {
    pub full_url: String,
    pub name: String,
}
impl DBTable for SearchURL {
    const TABLE_NAME: &'static str = "SearchURLs";
    const TABLE_SCHEMA: &'static str = "
    id           INTEGER PRIMARY KEY,
    url          TEXT NOT NULL UNIQUE,
    name         TEXT NOT NULL UNIQUE";

    fn get_all(conn: &Connection) -> rusqlite::Result<Vec<Self>> {
        let mut stmt = conn.prepare(&format!("SELECT url, name FROM {}", Self::TABLE_NAME))?;
        let iter = stmt.query_map([], |row| {
            Ok(SearchURL {
                full_url: row.get(0)?,
                name: row.get(1)?,
            })
        })?;

        let mut result = Vec::new();
        for item in iter {
            result.push(item?);
        }
        Ok(result)
    }
}
impl SearchURL {
    pub fn lookup(conn: &Connection, name: &str) -> Option<Self> {
        let mut stmt = conn
            .prepare(&format!(
                "SELECT * FROM {} WHERE name = ?",
                Self::TABLE_NAME
            ))
            .ok()?;
        stmt.query_one([name], |row| {
            Ok(SearchURL {
                // id: row.get(0)?,
                full_url: row.get(1)?,
                name: row.get(2)?,
            })
        })
        .ok()
    }

    pub fn add_or_update(&self, conn: &Connection) {
        let _ = conn
            .execute(
                &format!(
                    "INSERT OR REPLACE INTO {} (name, url) VALUES (?1, ?2)",
                    Self::TABLE_NAME
                ),
                (&self.name, &self.full_url),
            )
            .unwrap();
    }

    pub fn names(conn: &Connection) -> Vec<String> {
        let mut stmt = conn
            .prepare(&format!("SELECT name FROM {}", Self::TABLE_NAME))
            .ok()
            .unwrap();
        stmt.query_map([], |row| Ok(row.get(0)))
            .ok()
            .unwrap()
            .map(|e| e.unwrap())
            .flatten()
            .collect()
    }
}

#[derive(Serialize, Debug, PartialEq, Clone)]
pub struct ParsedPage {
    pub timestamp: DateTime<Utc>,
    pub category: String,
}
impl DBTable for ParsedPage {
    const TABLE_NAME: &'static str = "Pages_Parsed";
    const TABLE_SCHEMA: &'static str = "
        id                INTEGER PRIMARY KEY,
        category          TEXT NOT NULL,
        timestamp         INTEGER NOT NULL,
        UNIQUE(category, timestamp)
        FOREIGN KEY(category) REFERENCES SearchURLs(name)
    ";

    fn get_all(conn: &Connection) -> rusqlite::Result<Vec<Self>> {
        let mut stmt = conn.prepare(&format!(
            "SELECT category, timestamp FROM {}",
            Self::TABLE_NAME
        ))?;
        let iter = stmt.query_map([], |row| {
            Ok(ParsedPage {
                category: row.get(0)?,
                timestamp: row.get(1)?,
            })
        })?;

        let mut result = Vec::new();
        for item in iter {
            result.push(item?);
        }
        Ok(result)
    }
}
impl ParsedPage {
    pub fn lookup(conn: &Connection, timestamp: DateTime<Utc>) -> Option<Self> {
        let mut stmt = conn
            .prepare(&format!(
                "SELECT * FROM {} WHERE timestamp = ?",
                Self::TABLE_NAME
            ))
            .ok()?;
        stmt.query_one([timestamp], |row| {
            Ok(ParsedPage {
                // id: row.get(0)?,
                category: row.get(1)?,
                timestamp: row.get(2)?,
            })
        })
        .ok()
    }

    pub fn add_or_update(&self, conn: &Connection) {
        let _ = conn
            .execute(
                &format!(
                    "INSERT OR REPLACE INTO {} (category, timestamp) VALUES (?1, ?2)",
                    Self::TABLE_NAME
                ),
                (&self.category, self.timestamp),
            )
            .unwrap();
    }
}

#[derive(Serialize, Debug, PartialEq, Copy, Clone)]
pub struct ParsedStorage {
    pub id: i64,
    pub item: i64,
    pub total_gigabytes: i64,
    pub quantity: i64,
    pub individual_size_gigabytes: i64,
    pub parse_engine: i64,
    pub needed_description_check: bool,
}
impl DBTable for ParsedStorage {
    const TABLE_NAME: &'static str = "Storage_Parsed";
    const TABLE_SCHEMA: &'static str = "
        id                     INTEGER PRIMARY KEY,
        item                   INTEGER,
        total_gigabytes        INTEGER,
        quantity               INTEGER,
        sizes_gigabytes        TEXT,
        parse_engine           INTEGER,
        need_description_check INTEGER,
        UNIQUE(item, parse_engine)
        FOREIGN KEY(item) REFERENCES Listings(item_id)
    ";

    fn get_all(conn: &Connection) -> rusqlite::Result<Vec<Self>> {
        let mut stmt = conn.prepare(&format!("SELECT id, item, total_gigabytes, quantity, sizes_gigabytes, parse_engine, need_description_check FROM {}", Self::TABLE_NAME))?;
        let iter = stmt.query_map([], |row| {
            Ok(ParsedStorage {
                id: row.get(0)?,
                item: row.get(1)?,
                total_gigabytes: row.get(2)?,
                quantity: row.get(3)?,
                individual_size_gigabytes: {
                    let r: String = row.get(4)?;
                    r.parse().unwrap_or(0)
                },
                parse_engine: row.get(5)?,
                needed_description_check: row.get(6)?,
            })
        })?;

        let mut result = Vec::new();
        for item in iter {
            result.push(item?);
        }
        Ok(result)
    }
}
impl ParsedStorage {
    pub fn lookup(conn: &Connection, item: i64) -> Vec<ParsedStorage> {
        let mut stmt = conn
            .prepare(&format!("SELECT * FROM {} WHERE id = ?", Self::TABLE_NAME))
            .ok()
            .unwrap();
        stmt.query_map([item], |row| {
            Ok(ParsedStorage {
                id: row.get(0)?,
                item: row.get(1)?,
                total_gigabytes: row.get(2)?,
                quantity: row.get(3)?,
                individual_size_gigabytes: {
                    let r: String = row.get(4)?;
                    r.parse().unwrap()
                },
                parse_engine: row.get(5)?,
                needed_description_check: row.get(6)?,
            })
        })
        .ok()
        .unwrap()
        .map(|e| e.unwrap())
        .collect()
    }

    pub fn add_or_update(&self, conn: &Connection) {
        let _ = conn.execute(&format!("
            INSERT OR REPLACE INTO {}
                (item, total_gigabytes, quantity, sizes_gigabytes, parse_engine, need_description_check)
            VALUES
                (?1, ?2, ?3, ?4, ?5, ?6)",
                Self::TABLE_NAME),
            (
                &self.item,
                self.total_gigabytes,
                self.quantity,
                self.individual_size_gigabytes.to_string(),
                self.parse_engine,
                self.needed_description_check
            )
        ).unwrap();
    }
}

#[derive(Serialize, Debug, PartialEq, Clone)]
pub struct ItemAppearances {
    pub item: i64,
    pub timestamp: DateTime<Utc>,
    pub category: String,
    pub current_bid_usd_cents: Option<i64>,
}
impl DBTable for ItemAppearances {
    const TABLE_NAME: &'static str = "Item_Appearances";
    const TABLE_SCHEMA: &'static str = "
        id                      INTEGER PRIMARY KEY,
        item                    INTEGER NOT NULL,
        category                TEXT NOT NULL,
        timestamp               INTEGER NOT NULL,
        current_bid_usd_cents   INTEGER,
        UNIQUE(item, timestamp),
        FOREIGN KEY(item)       REFERENCES Listings(item_id),
        FOREIGN KEY(category, timestamp) REFERENCES Pages_Parsed(category, timestamp)
    ";

    fn get_all(conn: &Connection) -> rusqlite::Result<Vec<Self>> {
        let mut stmt = conn.prepare(&format!(
            "SELECT item, category, timestamp, current_bid_usd_cents FROM {}",
            Self::TABLE_NAME
        ))?;
        let iter = stmt.query_map([], |row| {
            Ok(ItemAppearances {
                item: row.get(0)?,
                category: row.get(1)?,
                timestamp: row.get(2)?,
                current_bid_usd_cents: row.get(3)?,
            })
        })?;

        let mut result = Vec::new();
        for item in iter {
            result.push(item?);
        }
        Ok(result)
    }
}
impl ItemAppearances {
    pub fn add_or_update(&self, conn: &Connection) {
        let count = conn
            .execute(
                &format!(
                    "
                    INSERT OR REPLACE INTO {}
                        (item, timestamp, category, current_bid_usd_cents)
                    VALUES
                        (?1, ?2, ?3, ?4)",
                    Self::TABLE_NAME
                ),
                (
                    self.item,
                    &self.timestamp,
                    &self.category,
                    self.current_bid_usd_cents,
                ),
            )
            .unwrap();
        if count != 1 {
            panic!("Expected count to be 1 but got {}", count);
        }
    }

    pub fn lookup(conn: &Connection, listing_id: i64) -> Vec<ItemAppearances> {
        let mut stmt = conn
            .prepare(&format!(
                "
                SELECT * FROM {}
                WHERE item IS ?1",
                Self::TABLE_NAME,
            ))
            .ok()
            .unwrap();
        stmt.query_map([listing_id], |row| {
            Ok(ItemAppearances {
                item: row.get(1)?,
                category: row.get(2)?,
                timestamp: row.get(3)?,
                current_bid_usd_cents: row.get(4)?,
            })
        })
        .ok()
        .unwrap()
        .map(|e| e.unwrap())
        .collect()
    }
}

#[derive(Serialize, Debug, PartialEq, Clone)]
pub struct Listing {
    pub id: i64,
    pub item_id: i64,
    pub title: String,
    pub buy_it_now_price_cents: Option<i64>,
    pub has_best_offer: bool,
    pub image_url: String,
}
impl DBTable for Listing {
    const TABLE_NAME: &'static str = "Listings";
    const TABLE_SCHEMA: &'static str = "
        id                     INTEGER PRIMARY KEY,
        item_id                INTEGER NOT NULL UNIQUE,
        title                  TEXT NOT NULL,
        buy_it_now_usd_cents   INTEGER,
        has_best_offer         INTEGER NOT NULL,
        image_url              TEXT NOT NULL
    ";

    fn get_all(conn: &Connection) -> rusqlite::Result<Vec<Self>> {
        let mut stmt = conn.prepare(&format!(
            "SELECT id, item_id, title, buy_it_now_usd_cents, has_best_offer, image_url FROM {}",
            Self::TABLE_NAME
        ))?;
        let iter = stmt.query_map([], |row| {
            Ok(Listing {
                id: row.get(0)?,
                item_id: row.get(1)?,
                title: row.get(2)?,
                buy_it_now_price_cents: row.get(3)?,
                has_best_offer: row.get(4)?,
                image_url: row.get(5)?,
            })
        })?;

        let mut result = Vec::new();
        for item in iter {
            result.push(item?);
        }
        Ok(result)
    }
}
impl Listing {
    pub fn lookup(conn: &Connection, item_id: i64) -> Option<Listing> {
        let mut stmt = conn
            .prepare(&format!(
                "SELECT * FROM {} WHERE item_id = ?",
                Self::TABLE_NAME
            ))
            .ok()?;
        stmt.query_one([item_id], |row| {
            Ok(Listing {
                id: row.get(0)?,
                item_id: row.get(1)?,
                title: row.get(2)?,
                buy_it_now_price_cents: row.get(3)?,
                has_best_offer: row.get(4)?,
                image_url: row.get(5)?,
            })
        })
        .ok()
    }

    pub fn lookup_since(conn: &Connection, since: DateTime<Utc>, limit: i64) -> Vec<Self> {
        let mut stmt = conn
            .prepare(&format!(
                "
                SELECT *
                FROM {0}
                WHERE EXISTS (
                  SELECT 1
                  FROM {1}
                  WHERE
                    {1}.item = {0}.item_id AND
                    {1}.timestamp >= ?1
                )
                LIMIT ?2
            ",
                Self::TABLE_NAME,
                ItemAppearances::TABLE_NAME
            ))
            .ok()
            .unwrap();

        stmt.query_map([since.timestamp(), limit], |row| {
            Ok(Listing {
                id: row.get(0)?,
                item_id: row.get(1)?,
                title: row.get(2)?,
                buy_it_now_price_cents: row.get(3)?,
                has_best_offer: row.get(4)?,
                image_url: row.get(5)?,
            })
        })
        .ok()
        .unwrap()
        .map(|e| e.unwrap())
        .collect()
    }

    pub fn lookup_non_parsed(conn: &Connection) -> Vec<(i64, String)> {
        let mut stmt = conn
            .prepare(&format!(
                "
                SELECT ei.item_id, ei.title FROM {} AS ei
                LEFT JOIN {} AS sp ON ei.item_id = sp.item
                WHERE sp.item IS NULL",
                Self::TABLE_NAME,
                ParsedStorage::TABLE_NAME
            ))
            .ok()
            .unwrap();
        stmt.query_map([], |row| Ok((row.get(0)?, row.get(1)?)))
            .ok()
            .unwrap()
            .map(|e| e.unwrap())
            .collect()
    }

    pub fn add_or_update(&self, conn: &Connection) {
        let count = conn
            .execute(
                &format!(
                    "INSERT OR REPLACE INTO {}
                (
                    item_id,
                    title,
                    buy_it_now_usd_cents,
                    has_best_offer,
                    image_url
                )
                VALUES (?1, ?2, ?3, ?4, ?5)",
                    Self::TABLE_NAME
                ),
                (
                    self.item_id,
                    &self.title,
                    self.buy_it_now_price_cents,
                    self.has_best_offer,
                    self.image_url.clone(),
                ),
            )
            .unwrap();
        if count != 1 {
            panic!("Expected count to be 1 but got {}", count);
        }
    }
}

pub fn get_initialized(path: Option<&Path>) -> Connection {
    let conn = match path {
        Some(p) => Connection::open(&p),
        None => Connection::open_in_memory(),
    }
    .unwrap();

    SearchURL::initialize(&conn);
    Listing::initialize(&conn);
    ParsedStorage::initialize(&conn);
    ParsedPage::initialize(&conn);
    ItemAppearances::initialize(&conn);

    conn
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn sanity_check() {
        let db = get_initialized(None);

        let searchurl = SearchURL {
            full_url: "google".to_owned(),
            name: "ssd".to_owned(),
        };
        searchurl.add_or_update(&db);
        assert_eq!(SearchURL::lookup(&db, &searchurl.name), Some(searchurl));

        let listing = Listing {
            id: 1,
            item_id: 1234,
            title: "Some Title".to_string(),
            buy_it_now_price_cents: Some(123),
            has_best_offer: false,
            image_url: "google.com".to_string(),
        };
        listing.add_or_update(&db);
        assert_eq!(Listing::lookup(&db, listing.item_id), Some(listing.clone()));

        let parsed = ParsedStorage {
            id: 1,
            item: 1234,
            total_gigabytes: 13,
            quantity: 3,
            individual_size_gigabytes: 13,
            parse_engine: 9,
            needed_description_check: true,
        };
        parsed.add_or_update(&db);
        assert_eq!(ParsedStorage::lookup(&db, listing.id), vec![parsed]);

        let page = ParsedPage {
            category: "ssd".to_owned(),
            timestamp: std::time::SystemTime::now().into(),
        };
        page.add_or_update(&db);
        assert_eq!(ParsedPage::lookup(&db, page.timestamp), Some(page.clone()));

        let apperance = ItemAppearances {
            item: listing.item_id,
            timestamp: page.timestamp,
            category: page.category,
            current_bid_usd_cents: Some(1233),
        };
        apperance.add_or_update(&db);
        assert_eq!(
            ItemAppearances::lookup(&db, listing.item_id),
            vec![apperance]
        );

        assert_eq!(Listing::lookup_since(&db, page.timestamp, 3), vec![listing]);
    }
}