use sqlx::Postgres;

use sqlx::Transaction;

// The flow for this function:

// 1. Create collection if it does not exists

// 2. Create the pipeline if it does not exist and add it to the collection.pipelines table with ACTIVE = FALSE

// 3. Create the tables for the collection_pipeline schema

// 4. Start a transaction

// 5. Sync the pipeline

// 6. Set the pipeline ACTIVE = TRUE

// 7. Commit the transaction

pipeline.verify_in_database(false).await?;

pipeline.create_tables().await?;

let pool = get_or_initialize_pool(&self.database_url).await?;

let transaction = pool.begin().await?;

let transaction = Arc::new(Mutex::new(transaction));

pipeline.execute(None, transaction.clone()).await?;

let mut transaction = Arc::into_inner(transaction)

.context("Error transaction dangling")?

.into_inner();

sqlx::query(&query_builder!(

"UPDATE %s SET active = TRUE WHERE name = $1",

self.pipelines_table_name

))

.bind(&pipeline.name)

.execute(&mut *transaction)

.await?;

transaction.commit().await?;

mp.println(format!("Done Syncing {}\n", pipeline.name))?;

pub async fn remove_pipeline(&mut self, pipeline: &MultiFieldPipeline) -> anyhow::Result<()> {

// The flow for this function:

// Create collection if it does not exist

// Begin a transaction

// Drop the collection_pipeline schema

// Delete the pipeline from the collection.pipelines table

// Commit the transaction

.context("Database data must be set to remove a pipeline from a collection")?

let pool = get_or_initialize_pool(&self.database_url).await?;

"DELETE FROM %s WHERE name = $1",

/// * `pipeline` - The [Pipeline] to enable

pub async fn enable_pipeline(

&mut self,

pipeline: &mut MultiFieldPipeline,

) -> anyhow::Result<()> {

self.add_pipeline(pipeline).await

/// * `pipeline` - The [Pipeline] to disable

pub async fn disable_pipeline(&mut self, pipeline: &MultiFieldPipeline) -> anyhow::Result<()> {

// Our current system for keeping documents, chunks, embeddings, and tsvectors in sync

// does not play nice with disabling and then re-enabling pipelines.

// For now, when disabling a pipeline, simply delete its schema and remake it later

// The flow for this function:

// 1. Create the collection if it does not exist

// 2. Begin a transaction

// 3. Set the pipelines ACTIVE = FALSE in the collection.pipelines table

// 4. Drop the collection_pipeline schema (this will get remade if they enable it again)

// 5. Commit the transaction

self.verify_in_database(false).await?;

let project_info = &self

.database_data

.as_ref()

.context("Database data must be set to remove a pipeline from a collection")?

.project_info;

let pool = get_or_initialize_pool(&self.database_url).await?;

let pipeline_schema = format!("{}_{}", project_info.name, pipeline.name);

let mut transaction = pool.begin().await?;

.execute(&mut *transaction)

transaction

.execute(query_builder!("DROP SCHEMA IF EXISTS %s CASCADE", pipeline_schema).as_str())

.await?;

transaction.commit().await?;

// The flow for this function

// 1. Create the collection if it does not exist

// 2. Get all pipelines where ACTIVE = TRUE

// 3. Create each pipeline and the collection_pipeline schema and tables if they don't already exist

// 4. Foreach document

// -> Begin a transaction returning the old document if it existed

// -> Insert the document

// -> Foreach pipeline check if we need to resync the document and if so sync the document

// -> Commit the transaction

let (document_id, previous_document): (i64, Option<Json>) = sqlx::query_as(&query_builder!(

"WITH prev AS (SELECT document FROM %s WHERE source_uuid = $1) INSERT INTO %s (source_uuid, document) VALUES ($1, $2) ON CONFLICT (source_uuid) DO UPDATE SET document = EXCLUDED.document RETURNING id, (SELECT document FROM prev)",

self.documents_table_name,

self.documents_table_name

))

.bind(&source_uuid)

.bind(&document)

.fetch_one(&mut *transaction)

.await?;

.map(|pipeline| {

(

pipeline,

previous_document.clone(),

document.clone(),

transaction.clone(),

)

.for_each_concurrent(

10,

|(pipeline, previous_document, document, transaction)| async move {

// Can unwrap here as we know it has parsed schema from the create_table call above

match previous_document {

Some(previous_document) => {

let should_run =

pipeline.parsed_schema.as_ref().unwrap().iter().any(

|(key, _)| document[key] != previous_document[key],

);

if should_run {

pipeline

.execute(Some(document_id), transaction)

.await

.expect("Failed to execute pipeline");

}

}

None => {

pipeline

.execute(Some(document_id), transaction)

.await

.expect("Failed to execute pipeline");

}

}

},

)

// #[instrument(skip(self))]

// async fn sync_pipeline(

// &mut self,

// pipeline: &mut MultiFieldPipeline,

// transaction: Arc<Mutex<Transaction<'static, Postgres>>>,

// ) -> anyhow::Result<()> {

// self.verify_in_database(false).await?;

// let project_info = &self

// .database_data

// .as_ref()

// .context("Database data must be set to get collection pipelines")?

// .project_info;

// pipeline.set_project_info(project_info.clone());

// pipeline.create_tables().await?;

// pipeline.execute(None, transaction).await?;

// Arc::into_inner(transaction)

// .context("Error transaction dangling")?

// .into_inner()

// .commit()

// .await?;

// Ok(())

// }

let collection_name = "test_r_c_capaud_46";

collection.archive().await?;

#[sqlx::test]

async fn can_update_documents() -> anyhow::Result<()> {

let collection_name = "test_r_c_cud_0";

let mut collection = Collection::new(collection_name, None);

let mut documents = generate_dummy_documents(1);

collection.upsert_documents(documents.clone(), None).await?;

documents[0]["body"] = json!("new body");

collection.upsert_documents(documents, None).await?;

// collection.archive().await?;

Ok(())

}