rust-web/src/pg_store.rs

use std::vec;

/// Store is responsible for manageing the in-memory hashmap of questions by providing initialization read/write functions,
/// and file I/O operations to persist these questions
/// TODO - Results returning errors should use specified types, not strings
use self::{
    answer::{Answer, NewAnswer},
    question::{NewQuestion, Question, QuestionDTO},
    question_tag::QuestionTag,
    tag::Tag,
};
use crate::*;

#[derive(Debug)]
pub struct Store {
    pub connection: PgPool,
}

impl Store {
    /// Helper to deal with unwrapping postgres fields as u8 which we want. Not gracefull if you pass in the wrong field name - so do not do that
    fn id_to_u8(pg_row: &PgRow, id_name: &str) -> u8 {
        pg_row.try_get::<i32, _>(id_name).unwrap() as u8
    }

    pub async fn new(db_url: &str) -> Self {
        let db_pool: Pool<Postgres> = match PgPoolOptions::new()
            .max_connections(5)
            .connect(db_url)
            .await
        {
            Ok(pool) => pool,
            Err(_) => panic!("Could not establish database connection"),
        };
        sqlx::migrate!().run(&db_pool).await.unwrap();
        Store {
            connection: db_pool,
        }
    }

    //Get the question tags for a particular question
    pub async fn get_tags_for_question(&self, question_id: u8) -> Option<Vec<Tag>> {
        let query = "SELECT tags.id, tags.label FROM tags 
        JOIN question_tag on tags.id = question_tag.tag_id 
        JOIN questions on question_tag.question_id = questions.id 
        WHERE questions.id = ($1)";
        let result = sqlx::query(query)
            .bind(i32::from(question_id))
            .fetch_all(&self.connection)
            .await;
        match result {
            Ok(pg_rows) => {
                let tags: Vec<Tag> = pg_rows
                    .iter()
                    .map(|pg_row| Tag::new(Store::id_to_u8(pg_row, "id"), pg_row.get("label")))
                    .collect::<Vec<Tag>>();
                Some(tags)
            }
            Err(e) => {
                println!("{}", e);
                None
            }
        }
    }

    //Remove a question/tag association
    pub async fn _remove_question_tag(
        &mut self,
        question_id: u8,
        tag_id: u8,
    ) -> Result<bool, String> {
        let result = sqlx::query("DELETE FROM question_tag WHERE question_id = $1 AND tag_id = $2")
            .bind(question_id.to_string())
            .bind(tag_id.to_string())
            .execute(&self.connection)
            .await;
        match result {
            Ok(_) => Ok(true),
            Err(e) => Err(e.to_string()),
        }
    }

    //Add a question/tag association
    pub async fn _add_question_tag(
        &mut self,
        question_id: u8,
        tag_id: u8,
    ) -> Result<QuestionTag, String> {
        let result =
            sqlx::query("INSERT INTO question_tag VALUES ($1,$2) RETURNING question_id, tag_id")
                .bind(question_id.to_string())
                .bind(tag_id.to_string())
                .fetch_one(&self.connection)
                .await;
        match result {
            Ok(pg_row) => Ok(QuestionTag::new(
                Store::id_to_u8(&pg_row, "question_id"),
                Store::id_to_u8(&pg_row, "tag_id"),
            )),
            Err(e) => Err(e.to_string()),
        }
    }

    // Add new tags to tags table, only if tags with existing label do not exist. If they already exist, just ignore that
    // Returns list of tags
    pub async fn add_tags(&mut self, tag_labels: Vec<String>) -> Result<Vec<Tag>, String> {
        let insert_query = "
        INSERT INTO tags (label)
        SELECT * FROM UNNEST(($1)::text[]) AS label
        ON CONFLICT (label) DO NOTHING
        ";

        //First run the insert query on the new labels
        let insert_result = sqlx::query(insert_query)
            .bind(&tag_labels)
            .execute(&self.connection)
            .await;
        match insert_result {
            Ok(_) => {
                //Then run the select query on the new labels (which may include already existing labels)
                let select_query = "
                    SELECT id, label FROM tags
                    WHERE label = ANY($1::text[]);
                ";
                match sqlx::query(select_query)
                    .bind(&tag_labels)
                    .fetch_all(&self.connection)
                    .await
                {
                    Ok(pg_rows) => {
                        let tags: Vec<Tag> = pg_rows
                            .iter()
                            .map(|pg_row| {
                                Tag::new(Store::id_to_u8(pg_row, "id"), pg_row.get("label"))
                            })
                            .collect();
                        Ok(tags)
                    }
                    Err(e) => Err(e.to_string()),
                }
            }
            Err(e) => Err(e.to_string()),
        }
    }

    /// Takes a question id and list of tags, and creates question_tag join associations
    /// Ignores if association already exists
    /// Returns list of the question tag associations
    pub async fn associate_tags(
        &mut self,
        question_id: u8,
        tags: &[Tag],
    ) -> Result<Vec<QuestionTag>, String> {
        let tag_ids_string: Vec<String> = tags.iter().map(|tag| tag.id.to_string()).collect();

        let query = "
        INSERT INTO question_tag (question_id, tag_id)
        SELECT $1::smallint, UNNEST($2::smallint[])
        ON CONFLICT DO NOTHING;
        ";
        let result = sqlx::query(query)
            .bind(question_id.to_string())
            .bind(tag_ids_string)
            .fetch_all(&self.connection)
            .await;
        match result {
            Ok(pg_rows) => {
                let question_tags: Vec<QuestionTag> = pg_rows
                    .iter()
                    .map(|pg_row| {
                        QuestionTag::new(
                            Store::id_to_u8(pg_row, "question_id"),
                            Store::id_to_u8(pg_row, "tag_id"),
                        )
                    })
                    .collect();
                Ok(question_tags)
            }
            Err(e) => Err(e.to_string()),
        }
    }

    // Takes a question id and a list of tags, and remove any existing question_tag join associations
    // Ignores if an association does not already exist
    // Returns Ok(true) on success
    pub async fn unassociate_tags(
        &mut self,
        question_id: u8,
        tags: Vec<Tag>,
    ) -> Result<bool, String> {
        let list = tags
            .iter()
            .map(|tag| format!("({},{})", question_id, tag.id))
            .collect::<Vec<String>>()
            .join(",");

        let query = format!(
            "DELETE FROM question_tag WHERE (question_id, tag_id) IN ({})",
            list
        ); // Not bulletproof to injection (still ok), but best we can do with sqlx aside from sequential hits
        match sqlx::query(&query).execute(&self.connection).await {
            Ok(_) => Ok(true),
            Err(e) => Err(e.to_string()),
        }
    }

    // Fetch a list of tags by either the tag id, or by the label. Up to the caller
    pub async fn _fetch_tags_by_property(
        &self,
        propert_id: u8,
        property_type: &str,
    ) -> Option<Vec<Tag>> {
        let query = format!("SELECT * FROM tags WHERE {} = ($2);", property_type).to_string(); //looks risky, but user does not get to control property type
        let result = sqlx::query(&query)
            .bind(property_type.to_string())
            .bind(propert_id.to_string())
            .fetch_all(&self.connection)
            .await;
        match result {
            Ok(pg_rows) => {
                let tags: Vec<Tag> = pg_rows
                    .iter()
                    .map(|pg_row| Tag::new(Store::id_to_u8(pg_row, "id"), pg_row.get("label")))
                    .collect();
                Some(tags)
            }
            Err(e) => {
                println!("err {}", e);
                None
            }
        }
    }

    //Add a new question - but also create the tags and question/tag associations as needed
    pub async fn add_question(&mut self, new_question: NewQuestion) -> Result<QuestionDTO, String> {
        let insert_question_result: Result<PgRow, sqlx::Error> = sqlx::query(
            "INSERT INTO questions (title, content) VALUES ($1, $2) RETURNING id, title, content",
        )
        .bind(new_question.title)
        .bind(new_question.content)
        .fetch_one(&self.connection)
        .await;
        // first, insert the new question
        match insert_question_result {
            Ok(pg_row) => {
                let inserted_question = Question::new(
                    Store::id_to_u8(&pg_row, "id"),
                    pg_row.get("title"),
                    pg_row.get("content"),
                );
                // Then create the new tags (if needed)
                let new_tags_result = &self.add_tags(new_question.tags).await;
                match new_tags_result {
                    Ok(new_tags_result) => {
                        let association = self
                            .associate_tags(inserted_question.id, new_tags_result)
                            .await;
                        // Finally, create the question/tag join
                        match association {
                            Ok(_) => {
                                // Now, return the QuestionDTO
                                Ok(QuestionDTO::new(
                                    inserted_question,
                                    new_tags_result.to_vec(),
                                ))
                            }
                            Err(e) => Err(e.to_string()),
                        }
                    }
                    Err(e) => Err(e.to_string()),
                }
            }
            Err(e) => Err(e.to_string()),
        }
    }

    // Delete the question tags association for the given question, then delete the question
    // Also clean up orphaned tags just in case
    pub async fn remove_question(&mut self, id: u8) -> Result<bool, String> {
        let delete_question_tag_query = "DELETE FROM question_tag WHERE question_id = $1";
        let delete_question_query = "DELETE FROM questions WHERE id = $1";
        // First, delete any possible question/tag associations
        match sqlx::query(delete_question_tag_query)
            .bind(i32::from(id))
            .execute(&self.connection)
            .await
        {
            Ok(_) => {
                //Now, delete the question
                match sqlx::query(delete_question_query)
                    .bind(i32::from(id))
                    .execute(&self.connection)
                    .await
                {
                    Ok(_) => {
                        // Finally remove any tags that may have been orphaned
                        match self.remove_orphan_tags().await {
                            Ok(_) => Ok(true),
                            Err(e) => Err(e.to_string()),
                        }
                    }
                    Err(e) => Err(e.to_string()),
                }
            }
            Err(e) => Err(e.to_string()),
        }
    }

    // Fetch one question, but do not worry about joining the tags
    pub async fn fetch_one_question_by_id(&self, id: u8) -> Option<Question> {
        let row_result = sqlx::query("SELECT id,title,content FROM questions WHERE id = $1")
            .bind(i32::from(id))
            .fetch_one(&self.connection)
            .await;
        match row_result {
            Ok(pg_row) => Some(Question::new(
                Store::id_to_u8(&pg_row, "id"),
                pg_row.get("title"),
                pg_row.get("content"),
            )),
            Err(e) => {
                println!("{}", e);
                None
            }
        }
    }

    // Fetch many questions - do not worry about joining the tags
    pub async fn fetch_many_questions(&self, start: i32, size: i32) -> Option<Vec<Question>> {
        let rows_result: Result<Vec<PgRow>, sqlx::Error> =
            sqlx::query("SELECT * FROM questions ORDER BY id LIMIT $1 OFFSET $2")
                .bind(size)
                .bind(start)
                .fetch_all(&self.connection)
                .await;

        match rows_result {
            Ok(pg_rows) => {
                let mut result: Vec<Question> = vec![];
                for pg_row in pg_rows {
                    result.push(Question::new(
                        Store::id_to_u8(&pg_row, "id"),
                        pg_row.get("title"),
                        pg_row.get("content"),
                    ))
                }
                Some(result)
            }
            Err(e) => {
                println!("{}", e);
                None
            }
        }
    }

    /// Remove tags from tags table, which have no question tag association
    /// Returns true on success
    pub async fn remove_orphan_tags(&mut self) -> Result<bool, String> {
        let result = sqlx::query(
            "DELETE FROM tags where id NOT IN (SELECT DISTINCT tag_id from question_tag)",
        )
        .execute(&self.connection)
        .await;
        match result {
            Ok(_) => Ok(true),
            Err(e) => Err(e.to_string()),
        }
    }

    //Update a question entity - just the question details (not tags)
    pub async fn update_question(
        &mut self,
        id: u8,
        title: String,
        content: String,
    ) -> Result<Question, String> {
        let query = "UPDATE questions
                    SET title = $1, content = $2 
                    WHERE id = $3 RETURNING id, title, content";
        let result = sqlx::query(query)
            .bind(title)
            .bind(content)
            .bind(i32::from(id))
            .fetch_one(&self.connection)
            .await;
        match result {
            Ok(pg_row) => Ok(Question::new(
                Store::id_to_u8(&pg_row, "id"),
                pg_row.get("title"),
                pg_row.get("content"),
            )),
            Err(e) => Err(e.to_string()),
        }
    }

    // Add an answer entity
    pub async fn add_answer(&mut self, new_answer: NewAnswer) -> Result<Answer, String> {
        let query = "INSERT INTO answers (content, question_id) VALUES ($1,$2) RETURNING id, content, question_id";
        let result = sqlx::query(query)
            .bind(new_answer.content)
            .bind(i32::from(new_answer.question_id))
            .fetch_one(&self.connection)
            .await;
        match result {
            Ok(pg_row) => Ok(Answer::new(
                Store::id_to_u8(&pg_row, "id"),
                pg_row.get("content"),
                Store::id_to_u8(&pg_row, "question_id"),
            )),
            Err(e) => Err(e.to_string()),
        }
    }
}