use crate::api::support_data_collector::SupportBundleData;
#[cfg(feature = "postgresql_embedded")]
use crate::config::PgEmbedConfig;
use crate::db::error::DBError;
use crate::db::operations;
#[cfg(feature = "postgresql_embedded")]
use crate::db::pg_setup;
use crate::db::storage::{ExtendedPipelineDescrRunner, Storage};
use crate::db::types::api_key::{ApiKeyDescr, ApiPermission};
use crate::db::types::monitor::{
    ClusterMonitorEvent, ClusterMonitorEventId, ExtendedClusterMonitorEvent, NewClusterMonitorEvent,
};
use crate::db::types::pipeline::{
    ExtendedPipelineDescr, ExtendedPipelineDescrMonitoring, PipelineDescr, PipelineId,
};
use crate::db::types::program::{
    ProgramConfig, ProgramInfo, ProgramStatus, RustCompilationInfo, SqlCompilationInfo,
};
use crate::db::types::resources_status::{ResourcesDesiredStatus, ResourcesStatus};
use crate::db::types::storage::StorageStatus;
use crate::db::types::tenant::TenantId;
use crate::db::types::version::Version;
use crate::is_supported_runtime;
use crate::{auth::TenantRecord, config::DatabaseConfig};
use async_trait::async_trait;
use deadpool_postgres::{Manager, Pool, RecyclingMethod};
use feldera_types::config::{PipelineConfig, RuntimeConfig};
use feldera_types::error::ErrorResponse;
use feldera_types::runtime_status::{BootstrapPolicy, ExtendedRuntimeStatus, RuntimeDesiredStatus};
use tokio_postgres::Row;
use tracing::{debug, info};
use uuid::Uuid;

// Convert PipelineId UUID to u64 to match PostgreSQL's behavior.
// This uses the first 8 bytes of the UUID converted to a big-endian u64.
// This ensures consistent worker assignment between Rust and SQL implementations.
fn pipline_uuid_to_u64(pipeline_id: PipelineId) -> u64 {
    let bytes = pipeline_id.0.as_bytes();
    u64::from_be_bytes([
        bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7],
    ])
}

/// Determine if a pipeline is assigned to a specific worker based on its ID.
/// Uses modulo operation to distribute pipelines across workers.
pub(crate) fn is_pipeline_assigned_to_worker(
    pipeline_id: PipelineId,
    worker_index: u64,
    total_workers: u64,
) -> bool {
    let hash = pipline_uuid_to_u64(pipeline_id);
    (hash % total_workers) == worker_index
}

mod embedded {
    use refinery::embed_migrations;
    embed_migrations!("./migrations/");
}

/// Postgres implementation of the `Storage` trait.
/// A trait method implementation consists of a single transaction
/// in which database operations are executed.
pub struct StoragePostgres {
    /// DatabaseConfig cmd line arguments.
    pub db_config: DatabaseConfig,
    /// Postgres configuration.
    pub config: tokio_postgres::Config,
    /// Pool from which clients can be taken.
    pub(crate) pool: Pool,
    /// Embedded Postgres instance which is used in development mode.
    /// It will not be persisted after termination.
    #[cfg(feature = "postgresql_embedded")]
    #[allow(dead_code)]
    // It has to stay alive until StoragePostgres is dropped.
    pub(crate) pg_inst: Option<postgresql_embedded::PostgreSQL>,
}

#[async_trait]
impl Storage for StoragePostgres {
    async fn check_connection(&self) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        operations::connectivity::check_connection(&txn).await?;
        txn.commit().await?;
        Ok(())
    }

    async fn get_or_create_tenant_id(
        &self,
        new_id: Uuid,
        name: String,
        provider: String,
    ) -> Result<TenantId, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let tenant_id =
            operations::tenant::get_or_create_tenant_id(&txn, new_id, name, provider).await?;
        txn.commit().await?;
        Ok(tenant_id)
    }

    async fn get_tenant_name(&self, tenant_id: TenantId) -> Result<String, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let tenant_name = operations::tenant::get_tenant_name(&txn, tenant_id).await?;
        txn.commit().await?;
        Ok(tenant_name)
    }

    async fn list_api_keys(&self, tenant_id: TenantId) -> Result<Vec<ApiKeyDescr>, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let api_keys = operations::api_key::list_api_keys(&txn, tenant_id).await?;
        txn.commit().await?;
        Ok(api_keys)
    }

    async fn get_api_key(&self, tenant_id: TenantId, name: &str) -> Result<ApiKeyDescr, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let api_key = operations::api_key::get_api_key(&txn, tenant_id, name).await?;
        txn.commit().await?;
        Ok(api_key)
    }

    async fn delete_api_key(&self, tenant_id: TenantId, name: &str) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let result = operations::api_key::delete_api_key(&txn, tenant_id, name).await?;
        txn.commit().await?;
        Ok(result)
    }

    async fn store_api_key_hash(
        &self,
        tenant_id: TenantId,
        id: Uuid,
        name: &str,
        key: &str,
        permissions: Vec<ApiPermission>,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let result =
            operations::api_key::store_api_key_hash(&txn, tenant_id, id, name, key, permissions)
                .await?;
        txn.commit().await?;
        Ok(result)
    }

    async fn validate_api_key(&self, key: &str) -> Result<(TenantId, Vec<ApiPermission>), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let result = operations::api_key::validate_api_key(&txn, key).await?;
        txn.commit().await?;
        Ok(result)
    }

    async fn list_pipelines(
        &self,
        tenant_id: TenantId,
    ) -> Result<Vec<ExtendedPipelineDescr>, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipelines = operations::pipeline::list_pipelines(&txn, tenant_id).await?;
        txn.commit().await?;
        Ok(pipelines)
    }

    async fn list_pipelines_for_monitoring(
        &self,
        tenant_id: TenantId,
    ) -> Result<Vec<ExtendedPipelineDescrMonitoring>, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipelines =
            operations::pipeline::list_pipelines_for_monitoring(&txn, tenant_id).await?;
        txn.commit().await?;
        Ok(pipelines)
    }

    async fn get_pipeline(
        &self,
        tenant_id: TenantId,
        name: &str,
    ) -> Result<ExtendedPipelineDescr, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipeline = operations::pipeline::get_pipeline(&txn, tenant_id, name).await?;
        txn.commit().await?;
        Ok(pipeline)
    }

    async fn get_pipeline_for_monitoring(
        &self,
        tenant_id: TenantId,
        name: &str,
    ) -> Result<ExtendedPipelineDescrMonitoring, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipeline =
            operations::pipeline::get_pipeline_for_monitoring(&txn, tenant_id, name).await?;
        txn.commit().await?;
        Ok(pipeline)
    }

    async fn get_pipeline_by_id(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
    ) -> Result<ExtendedPipelineDescr, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipeline =
            operations::pipeline::get_pipeline_by_id(&txn, tenant_id, pipeline_id).await?;
        txn.commit().await?;
        Ok(pipeline)
    }

    async fn get_pipeline_by_id_for_monitoring(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
    ) -> Result<ExtendedPipelineDescrMonitoring, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipeline =
            operations::pipeline::get_pipeline_by_id_for_monitoring(&txn, tenant_id, pipeline_id)
                .await?;
        txn.commit().await?;
        Ok(pipeline)
    }

    async fn get_pipeline_by_id_for_runner(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
        platform_version: &str,
        provision_called: bool,
    ) -> Result<ExtendedPipelineDescrRunner, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipeline_monitoring =
            operations::pipeline::get_pipeline_by_id_for_monitoring(&txn, tenant_id, pipeline_id)
                .await?;
        let is_ready_compiled = pipeline_monitoring.program_status == ProgramStatus::Success
            && is_supported_runtime(platform_version, &pipeline_monitoring.platform_version);
        let pipeline_result = if matches!(
            (
                pipeline_monitoring.deployment_resources_status,
                pipeline_monitoring.deployment_resources_desired_status,
                is_ready_compiled,
                provision_called
            ),
            (
                ResourcesStatus::Stopped,
                ResourcesDesiredStatus::Provisioned,
                true,
                _
            ) | (
                ResourcesStatus::Provisioning,
                ResourcesDesiredStatus::Provisioned,
                _,
                false
            ),
        ) {
            ExtendedPipelineDescrRunner::Complete(
                operations::pipeline::get_pipeline_by_id(&txn, tenant_id, pipeline_id).await?,
            )
        } else {
            ExtendedPipelineDescrRunner::Monitoring(pipeline_monitoring)
        };
        txn.commit().await?;
        Ok(pipeline_result)
    }

    async fn new_pipeline(
        &self,
        tenant_id: TenantId,
        new_id: Uuid,
        platform_version: &str,
        pipeline: PipelineDescr,
    ) -> Result<ExtendedPipelineDescr, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;

        // Create new pipeline
        operations::pipeline::new_pipeline(
            &txn,
            tenant_id,
            new_id,
            platform_version,
            pipeline.clone(),
        )
        .await?;

        // Fetch newly created pipeline
        let extended_pipeline =
            operations::pipeline::get_pipeline(&txn, tenant_id, &pipeline.name).await?;

        txn.commit().await?;
        Ok(extended_pipeline)
    }

    async fn new_or_update_pipeline(
        &self,
        tenant_id: TenantId,
        new_id: Uuid,
        original_name: &str,
        platform_version: &str,
        bump_platform_version: bool,
        pipeline: PipelineDescr,
    ) -> Result<(bool, ExtendedPipelineDescr), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;

        // Check if pipeline exists
        let current = operations::pipeline::get_pipeline(&txn, tenant_id, original_name).await;
        let is_new: bool = match current {
            Ok(_) => {
                // Pipeline already exists, as such update it
                operations::pipeline::update_pipeline(
                    &txn,
                    false, // Done by user
                    tenant_id,
                    original_name,
                    &Some(pipeline.name.clone()),
                    &Some(pipeline.description.clone()),
                    platform_version,
                    bump_platform_version,
                    &Some(pipeline.runtime_config.clone()),
                    &Some(pipeline.program_code.clone()),
                    &Some(pipeline.udf_rust.clone()),
                    &Some(pipeline.udf_toml.clone()),
                    &Some(pipeline.program_config.clone()),
                )
                .await?;
                false
            }
            Err(DBError::UnknownPipelineName { .. }) => {
                // Pipeline does not yet exist, as such create it
                if original_name != pipeline.name {
                    return Err(DBError::CannotRenameNonExistingPipeline);
                }
                operations::pipeline::new_pipeline(
                    &txn,
                    tenant_id,
                    new_id,
                    platform_version,
                    pipeline.clone(),
                )
                .await?;
                true
            }
            Err(e) => {
                // Another error occurred when fetching the extended pipeline,
                // return that if that is the case
                return Err(e);
            }
        };

        // Fetch new or updated pipeline
        let extended_pipeline =
            operations::pipeline::get_pipeline(&txn, tenant_id, &pipeline.name).await?;

        txn.commit().await?;
        Ok((is_new, extended_pipeline))
    }

    async fn testing_force_update_platform_version(
        &self,
        tenant_id: TenantId,
        pipeline_name: &str,
        platform_version: &str,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;

        // Check if pipeline exists
        let current =
            operations::pipeline::get_pipeline_for_monitoring(&txn, tenant_id, pipeline_name)
                .await?;

        if current.deployment_resources_status != ResourcesStatus::Stopped {
            return Err(DBError::UpdateRestrictedToStopped);
        }

        let stmt_update = txn
            .prepare_cached(
                "UPDATE pipeline
                     SET platform_version = $1
                     WHERE id = $2",
            )
            .await?;

        let rows_affected = txn
            .execute(&stmt_update, &[&platform_version, &current.id.0])
            .await?;

        assert_eq!(rows_affected, 1);

        info!(
            pipeline_id = %current.id,
            pipeline = %current.name,
            "Updated pipeline platform version from {} to {}",
            current.platform_version,
            platform_version
        );

        txn.commit().await?;

        Ok(())
    }

    #[allow(clippy::too_many_arguments)]
    async fn update_pipeline(
        &self,
        tenant_id: TenantId,
        original_name: &str,
        name: &Option<String>,
        description: &Option<String>,
        platform_version: &str,
        bump_platform_version: bool,
        runtime_config: &Option<serde_json::Value>,
        program_code: &Option<String>,
        udf_rust: &Option<String>,
        udf_toml: &Option<String>,
        program_config: &Option<serde_json::Value>,
    ) -> Result<ExtendedPipelineDescr, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;

        // Update existing pipeline
        operations::pipeline::update_pipeline(
            &txn,
            false, // Done by user
            tenant_id,
            original_name,
            name,
            description,
            platform_version,
            bump_platform_version,
            runtime_config,
            program_code,
            udf_rust,
            udf_toml,
            program_config,
        )
        .await?;

        // Fetch updated pipeline
        let final_name = name.clone().unwrap_or(original_name.to_string());
        let extended_pipeline =
            operations::pipeline::get_pipeline(&txn, tenant_id, &final_name).await?;

        txn.commit().await?;
        Ok(extended_pipeline)
    }

    async fn delete_pipeline(
        &self,
        tenant_id: TenantId,
        pipeline_name: &str,
    ) -> Result<PipelineId, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipeline_id =
            operations::pipeline::delete_pipeline(&txn, tenant_id, pipeline_name).await?;
        txn.commit().await?;
        Ok(pipeline_id)
    }

    async fn transit_program_status_to_pending(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
        program_version_guard: Version,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        operations::pipeline::set_program_status(
            &txn,
            tenant_id,
            pipeline_id,
            program_version_guard,
            &ProgramStatus::Pending,
            &None,
            &None,
            &None,
            &None,
            &None,
            &None,
            &None,
        )
        .await?;
        txn.commit().await?;
        Ok(())
    }

    async fn transit_program_status_to_compiling_sql(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
        program_version_guard: Version,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        operations::pipeline::set_program_status(
            &txn,
            tenant_id,
            pipeline_id,
            program_version_guard,
            &ProgramStatus::CompilingSql,
            &None,
            &None,
            &None,
            &None,
            &None,
            &None,
            &None,
        )
        .await?;
        txn.commit().await?;
        Ok(())
    }

    async fn transit_program_status_to_sql_compiled(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
        program_version_guard: Version,
        sql_compilation: &SqlCompilationInfo,
        program_info: &serde_json::Value,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        operations::pipeline::set_program_status(
            &txn,
            tenant_id,
            pipeline_id,
            program_version_guard,
            &ProgramStatus::SqlCompiled,
            &Some(sql_compilation.clone()),
            &None,
            &None,
            &Some(program_info.clone()),
            &None,
            &None,
            &None,
        )
        .await?;
        txn.commit().await?;
        Ok(())
    }

    async fn transit_program_status_to_compiling_rust(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
        program_version_guard: Version,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        operations::pipeline::set_program_status(
            &txn,
            tenant_id,
            pipeline_id,
            program_version_guard,
            &ProgramStatus::CompilingRust,
            &None,
            &None,
            &None,
            &None,
            &None,
            &None,
            &None,
        )
        .await?;
        txn.commit().await?;
        Ok(())
    }

    async fn transit_program_status_to_success(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
        program_version_guard: Version,
        rust_compilation: &RustCompilationInfo,
        program_binary_source_checksum: &str,
        program_binary_integrity_checksum: &str,
        program_info_integrity_checksum: &str,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        operations::pipeline::set_program_status(
            &txn,
            tenant_id,
            pipeline_id,
            program_version_guard,
            &ProgramStatus::Success,
            &None,
            &Some(rust_compilation.clone()),
            &None,
            &None,
            &Some(program_binary_source_checksum.to_string()),
            &Some(program_binary_integrity_checksum.to_string()),
            &Some(program_info_integrity_checksum.to_string()),
        )
        .await?;
        txn.commit().await?;
        Ok(())
    }

    async fn transit_program_status_to_sql_error(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
        program_version_guard: Version,
        sql_compilation: &SqlCompilationInfo,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        operations::pipeline::set_program_status(
            &txn,
            tenant_id,
            pipeline_id,
            program_version_guard,
            &ProgramStatus::SqlError,
            &Some(sql_compilation.clone()),
            &None,
            &None,
            &None,
            &None,
            &None,
            &None,
        )
        .await?;
        txn.commit().await?;
        Ok(())
    }

    async fn transit_program_status_to_rust_error(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
        program_version_guard: Version,
        rust_compilation: &RustCompilationInfo,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        operations::pipeline::set_program_status(
            &txn,
            tenant_id,
            pipeline_id,
            program_version_guard,
            &ProgramStatus::RustError,
            &None,
            &Some(rust_compilation.clone()),
            &None,
            &None,
            &None,
            &None,
            &None,
        )
        .await?;
        txn.commit().await?;
        Ok(())
    }

    async fn transit_program_status_to_system_error(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
        program_version_guard: Version,
        system_error: &str,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        operations::pipeline::set_program_status(
            &txn,
            tenant_id,
            pipeline_id,
            program_version_guard,
            &ProgramStatus::SystemError,
            &None,
            &None,
            &Some(system_error.to_string()),
            &None,
            &None,
            &None,
            &None,
        )
        .await?;
        txn.commit().await?;
        Ok(())
    }

    async fn dismiss_deployment_error(
        &self,
        tenant_id: TenantId,
        pipeline_name: &str,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        operations::pipeline::dismiss_deployment_error(&txn, tenant_id, pipeline_name).await?;
        txn.commit().await?;
        Ok(())
    }

    async fn set_deployment_resources_desired_status_provisioned(
        &self,
        tenant_id: TenantId,
        pipeline_name: &str,
        initial: RuntimeDesiredStatus,
        bootstrap_policy: BootstrapPolicy,
        dismiss_error: bool,
    ) -> Result<PipelineId, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipeline_id = operations::pipeline::set_deployment_resources_desired_status(
            &txn,
            tenant_id,
            pipeline_name,
            ResourcesDesiredStatus::Provisioned,
            Some(initial),
            Some(bootstrap_policy),
            dismiss_error,
        )
        .await?;
        txn.commit().await?;
        Ok(pipeline_id)
    }

    async fn set_deployment_resources_desired_status_stopped(
        &self,
        tenant_id: TenantId,
        pipeline_name: &str,
    ) -> Result<PipelineId, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipeline_id = operations::pipeline::set_deployment_resources_desired_status(
            &txn,
            tenant_id,
            pipeline_name,
            ResourcesDesiredStatus::Stopped,
            None,
            None,
            false,
        )
        .await?;
        txn.commit().await?;
        Ok(pipeline_id)
    }

    async fn set_deployment_resources_desired_status_stopped_if_not_provisioned(
        &self,
        tenant_id: TenantId,
        pipeline_name: &str,
    ) -> Result<(bool, PipelineId), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipeline =
            operations::pipeline::get_pipeline_for_monitoring(&txn, tenant_id, pipeline_name)
                .await?;
        let was_set = if pipeline.deployment_resources_status != ResourcesStatus::Provisioned {
            operations::pipeline::set_deployment_resources_desired_status(
                &txn,
                tenant_id,
                pipeline_name,
                ResourcesDesiredStatus::Stopped,
                None,
                None,
                false,
            )
            .await?;
            true
        } else {
            false
        };
        txn.commit().await?;
        Ok((was_set, pipeline.id))
    }

    async fn transit_deployment_resources_status_to_provisioning(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
        version_guard: Version,
        deployment_id: Uuid,
        deployment_config: serde_json::Value,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        // If the pipeline currently is already Stopped and is desired to be Stopped,
        // then there is no need to transition to Provisioning
        let pipeline =
            operations::pipeline::get_pipeline_by_id_for_monitoring(&txn, tenant_id, pipeline_id)
                .await?;
        if pipeline.deployment_resources_status == ResourcesStatus::Stopped
            && pipeline.deployment_resources_desired_status == ResourcesDesiredStatus::Stopped
        {
            return Err(DBError::UnnecessaryResourcesStatusTransition {
                current_status: pipeline.deployment_resources_status,
                new_status: ResourcesStatus::Provisioning,
                current_desired_status: pipeline.deployment_resources_desired_status,
            });
        }
        operations::pipeline::set_storage_status(
            &txn,
            tenant_id,
            pipeline_id,
            StorageStatus::InUse,
        )
        .await?;
        operations::pipeline::set_deployment_resources_status_provisioning(
            &txn,
            tenant_id,
            pipeline_id,
            version_guard,
            deployment_id,
            deployment_config,
        )
        .await?;
        txn.commit().await?;
        Ok(())
    }

    async fn remain_deployment_resources_status_provisioning(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
        version_guard: Version,
        deployment_resources_status_details: serde_json::Value,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        operations::pipeline::remain_deployment_resources_status_provisioning(
            &txn,
            tenant_id,
            pipeline_id,
            version_guard,
            deployment_resources_status_details,
        )
        .await?;
        txn.commit().await?;
        Ok(())
    }

    async fn transit_deployment_resources_status_to_provisioned(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
        version_guard: Version,
        deployment_location: &str,
        deployment_resources_status_details: serde_json::Value,
        extended_runtime_status: ExtendedRuntimeStatus,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        operations::pipeline::set_deployment_resources_status_provisioned(
            &txn,
            tenant_id,
            pipeline_id,
            version_guard,
            deployment_location.to_string(),
            deployment_resources_status_details,
            extended_runtime_status.runtime_status,
            extended_runtime_status.runtime_status_details,
            extended_runtime_status.runtime_desired_status,
        )
        .await?;
        txn.commit().await?;
        Ok(())
    }

    async fn remain_deployment_resources_status_provisioned(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
        version_guard: Version,
        deployment_resources_status_details: serde_json::Value,
        extended_runtime_status: ExtendedRuntimeStatus,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        operations::pipeline::remain_deployment_resources_status_provisioned(
            &txn,
            tenant_id,
            pipeline_id,
            version_guard,
            deployment_resources_status_details,
            extended_runtime_status.runtime_status,
            extended_runtime_status.runtime_status_details,
            extended_runtime_status.runtime_desired_status,
        )
        .await?;
        txn.commit().await?;
        Ok(())
    }

    async fn transit_deployment_resources_status_to_stopping(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
        version_guard: Version,
        deployment_error: Option<ErrorResponse>,
        suspend_info: Option<serde_json::Value>,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipeline =
            operations::pipeline::get_pipeline_by_id_for_monitoring(&txn, tenant_id, pipeline_id)
                .await?;
        // If the pipeline currently is already Stopped and is desired to be Stopped,
        // then there is no need to transition to Stopping
        if pipeline.deployment_resources_status == ResourcesStatus::Stopped
            && pipeline.deployment_resources_desired_status == ResourcesDesiredStatus::Stopped
        {
            return Err(DBError::UnnecessaryResourcesStatusTransition {
                current_status: pipeline.deployment_resources_status,
                new_status: ResourcesStatus::Stopping,
                current_desired_status: pipeline.deployment_resources_desired_status,
            });
        }
        operations::pipeline::set_deployment_resources_desired_status(
            &txn,
            tenant_id,
            &pipeline.name,
            ResourcesDesiredStatus::Stopped,
            None,
            None,
            false,
        )
        .await?;
        operations::pipeline::set_deployment_resources_status_stopping(
            &txn,
            tenant_id,
            pipeline_id,
            version_guard,
            deployment_error,
            suspend_info,
        )
        .await?;

        txn.commit().await?;
        Ok(())
    }

    async fn remain_deployment_resources_status_stopping(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
        version_guard: Version,
        deployment_resources_status_details: serde_json::Value,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        operations::pipeline::remain_deployment_resources_status_stopping(
            &txn,
            tenant_id,
            pipeline_id,
            version_guard,
            deployment_resources_status_details,
        )
        .await?;

        txn.commit().await?;
        Ok(())
    }

    async fn transit_deployment_resources_status_to_stopped(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
        version_guard: Version,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        operations::pipeline::set_deployment_resources_status_stopped(
            &txn,
            tenant_id,
            pipeline_id,
            version_guard,
        )
        .await?;
        txn.commit().await?;
        Ok(())
    }

    async fn transit_storage_status_to_clearing_if_not_cleared(
        &self,
        tenant_id: TenantId,
        pipeline_name: &str,
    ) -> Result<PipelineId, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipeline =
            operations::pipeline::get_pipeline_for_monitoring(&txn, tenant_id, pipeline_name)
                .await?;
        if pipeline.storage_status != StorageStatus::Cleared {
            // If it is already cleared, it does not have to transition to clearing
            operations::pipeline::set_storage_status(
                &txn,
                tenant_id,
                pipeline.id,
                StorageStatus::Clearing,
            )
            .await?;
        }
        txn.commit().await?;
        Ok(pipeline.id)
    }

    async fn transit_storage_status_to_cleared(
        &self,
        tenant_id: TenantId,
        pipeline_id: PipelineId,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        operations::pipeline::set_storage_status(
            &txn,
            tenant_id,
            pipeline_id,
            StorageStatus::Cleared,
        )
        .await?;
        txn.commit().await?;
        Ok(())
    }

    /// Increments notify counter such that any LISTEN mechanism is appraised.
    async fn increment_notify_counter(
        &self,
        tenant_id: TenantId,
        pipeline_name: &str,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        operations::pipeline::increment_notify_counter(&txn, tenant_id, pipeline_name).await?;
        txn.commit().await?;
        Ok(())
    }

    async fn list_pipeline_ids_across_all_tenants(
        &self,
    ) -> Result<Vec<(TenantId, PipelineId)>, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipeline_ids = operations::pipeline::list_pipeline_ids_across_all_tenants(&txn).await?;
        txn.commit().await?;
        Ok(pipeline_ids)
    }

    async fn list_pipelines_across_all_tenants_for_monitoring(
        &self,
    ) -> Result<Vec<(TenantId, ExtendedPipelineDescrMonitoring)>, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipelines =
            operations::pipeline::list_pipelines_across_all_tenants_for_monitoring(&txn).await?;
        txn.commit().await?;
        Ok(pipelines)
    }

    async fn clear_ongoing_sql_compilation_for_worker(
        &self,
        platform_version: &str,
        worker_id: usize,
        total_workers: usize,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipelines =
            operations::pipeline::list_pipelines_across_all_tenants_for_monitoring(&txn).await?;
        for (tenant_id, pipeline) in pipelines {
            // skip the pipelines that are not assigned to this worker
            if !is_pipeline_assigned_to_worker(pipeline.id, worker_id as u64, total_workers as u64)
            {
                continue;
            }

            if pipeline.deployment_resources_status == ResourcesStatus::Stopped {
                if pipeline.platform_version == platform_version {
                    if pipeline.program_status == ProgramStatus::CompilingSql {
                        operations::pipeline::set_program_status(
                            &txn,
                            tenant_id,
                            pipeline.id,
                            pipeline.program_version,
                            &ProgramStatus::Pending,
                            &None,
                            &None,
                            &None,
                            &None,
                            &None,
                            &None,
                            &None,
                        )
                        .await?;
                    }
                } else if pipeline.program_status == ProgramStatus::Pending
                    || pipeline.program_status == ProgramStatus::CompilingSql
                {
                    operations::pipeline::update_pipeline(
                        &txn,
                        true, // Done by compiler
                        tenant_id,
                        &pipeline.name,
                        &None,
                        &None,
                        platform_version,
                        true,
                        &None,
                        &None,
                        &None,
                        &None,
                        &None,
                    )
                    .await?;
                }
            }
        }
        txn.commit().await?;
        Ok(())
    }

    async fn get_next_sql_compilation(
        &self,
        platform_version: &str,
        worker_id: usize,
        total_workers: usize,
    ) -> Result<Option<(TenantId, ExtendedPipelineDescr)>, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let next_pipeline_program = operations::pipeline::get_next_sql_compilation(
            &txn,
            platform_version,
            worker_id,
            total_workers,
        )
        .await?;
        txn.commit().await?;
        Ok(next_pipeline_program)
    }

    async fn clear_ongoing_rust_compilation_for_worker(
        &self,
        platform_version: &str,
        worker_id: usize,
        total_workers: usize,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipelines =
            operations::pipeline::list_pipelines_across_all_tenants_for_monitoring(&txn).await?;

        for (tenant_id, pipeline) in pipelines {
            // skip the pipelines that are not assigned to this worker
            if !is_pipeline_assigned_to_worker(pipeline.id, worker_id as u64, total_workers as u64)
            {
                continue;
            }

            if pipeline.deployment_resources_status == ResourcesStatus::Stopped {
                if pipeline.platform_version == platform_version {
                    if pipeline.program_status == ProgramStatus::CompilingRust {
                        // Because `program_info` can be rather large, it is only fetched when
                        // the program status needs to be reset to `SqlCompiled`
                        let pipeline_complete =
                            operations::pipeline::get_pipeline_by_id(&txn, tenant_id, pipeline.id)
                                .await?;
                        operations::pipeline::set_program_status(
                            &txn,
                            tenant_id,
                            pipeline.id,
                            pipeline.program_version,
                            &ProgramStatus::SqlCompiled,
                            &Some(pipeline_complete.program_error.sql_compilation.clone().expect("program_error.sql_compilation must be present if current status is CompilingRust")),
                            &None,
                            &None,
                            &Some(pipeline_complete.program_info.clone().expect(
                                "program_info must be present if current status is CompilingRust",
                            )),
                            &None,
                            &None,
                            &None,
                        )
                        .await?;
                    }
                } else if pipeline.program_status == ProgramStatus::SqlCompiled
                    || pipeline.program_status == ProgramStatus::CompilingRust
                {
                    operations::pipeline::update_pipeline(
                        &txn,
                        true, // Done by compiler
                        tenant_id,
                        &pipeline.name,
                        &None,
                        &None,
                        platform_version,
                        true,
                        &None,
                        &None,
                        &None,
                        &None,
                        &None,
                    )
                    .await?;
                }
            }
        }
        txn.commit().await?;
        Ok(())
    }

    async fn get_next_rust_compilation(
        &self,
        platform_version: &str,
        worker_id: usize,
        total_workers: usize,
    ) -> Result<Option<(TenantId, ExtendedPipelineDescr)>, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let next_pipeline_program = operations::pipeline::get_next_rust_compilation(
            &txn,
            platform_version,
            worker_id,
            total_workers,
        )
        .await?;
        txn.commit().await?;
        Ok(next_pipeline_program)
    }

    async fn list_pipeline_programs_across_all_tenants(
        &self,
    ) -> Result<
        Vec<(
            PipelineId,
            Version,
            Option<String>,
            Option<String>,
            Option<String>,
        )>,
        DBError,
    > {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipeline_programs =
            operations::pipeline::list_pipeline_programs_across_all_tenants(&txn).await?;
        txn.commit().await?;
        Ok(pipeline_programs)
    }

    async fn get_support_bundle_data(
        &self,
        tenant_id: TenantId,
        pipeline_name: &str,
        how_many: u64,
    ) -> Result<(ExtendedPipelineDescrMonitoring, Vec<SupportBundleData>), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let pipeline =
            operations::pipeline::get_pipeline_for_monitoring(&txn, tenant_id, pipeline_name)
                .await?;
        let bundle_data =
            operations::pipeline::get_support_bundle_data(&txn, pipeline.id, how_many).await?;
        txn.commit().await?;
        Ok((pipeline, bundle_data))
    }

    async fn list_cluster_monitor_events(&self) -> Result<Vec<ClusterMonitorEvent>, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let events = operations::cluster_monitor::list_cluster_monitor_events_short(&txn).await?;
        txn.commit().await?;
        Ok(events)
    }

    async fn get_cluster_monitor_event_short(
        &self,
        event_id: ClusterMonitorEventId,
    ) -> Result<ClusterMonitorEvent, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let event =
            operations::cluster_monitor::get_cluster_monitor_event_short(&txn, event_id).await?;
        txn.commit().await?;
        Ok(event)
    }

    async fn get_cluster_monitor_event_extended(
        &self,
        event_id: ClusterMonitorEventId,
    ) -> Result<ExtendedClusterMonitorEvent, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let event =
            operations::cluster_monitor::get_cluster_monitor_event_extended(&txn, event_id).await?;
        txn.commit().await?;
        Ok(event)
    }

    async fn get_latest_cluster_monitor_event_short(&self) -> Result<ClusterMonitorEvent, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let event =
            operations::cluster_monitor::get_latest_cluster_monitor_event_short(&txn).await?;
        txn.commit().await?;
        Ok(event)
    }

    async fn get_latest_cluster_monitor_event_extended(
        &self,
    ) -> Result<ExtendedClusterMonitorEvent, DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let event =
            operations::cluster_monitor::get_latest_cluster_monitor_event_extended(&txn).await?;
        txn.commit().await?;
        Ok(event)
    }

    async fn new_cluster_monitor_event(
        &self,
        new_id: Uuid,
        event_descr: NewClusterMonitorEvent,
    ) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        operations::cluster_monitor::new_cluster_monitor_event(&txn, new_id, event_descr).await?;
        txn.commit().await?;
        Ok(())
    }

    async fn delete_cluster_monitor_events_beyond_retention(
        &self,
        retention_hours: u16,
        retention_num: u16,
    ) -> Result<(u64, u64), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let (num_deleted_due_to_timestamp, num_deleted_due_to_limit) =
            operations::cluster_monitor::delete_cluster_monitor_events_beyond_retention(
                &txn,
                retention_hours,
                retention_num,
            )
            .await?;
        txn.commit().await?;
        Ok((num_deleted_due_to_timestamp, num_deleted_due_to_limit))
    }
}

impl StoragePostgres {
    pub async fn connect(
        db_config: &DatabaseConfig,
        #[cfg(feature = "postgresql_embedded")] pg_embed_config: PgEmbedConfig,
    ) -> Result<Self, DBError> {
        #[cfg(feature = "postgresql_embedded")]
        if db_config.uses_postgres_embed() {
            let database_dir = pg_embed_config.pg_embed_data_dir();
            let pg_inst = pg_setup::install(database_dir, true, Some(8082)).await?;
            let db_name = "postgres";
            if !pg_inst.database_exists(db_name).await? {
                pg_inst.create_database(db_name).await?;
            }
            let connection_string = pg_inst.settings().url(db_name);

            let db_config = DatabaseConfig::new(connection_string, None);
            return Self::initialize(&db_config, Some(pg_inst)).await;
        };

        Self::initialize(
            db_config,
            #[cfg(feature = "postgresql_embedded")]
            None,
        )
        .await
    }

    pub(crate) async fn initialize(
        db_config: &DatabaseConfig,
        #[cfg(feature = "postgresql_embedded")] pg_inst: Option<postgresql_embedded::PostgreSQL>,
    ) -> Result<Self, DBError> {
        let config = db_config.tokio_postgres_config()?;
        debug!(
            "Opening Postgres connection with configuration: {:?}",
            config
        );
        let tls = db_config.tls_connector()?;
        let mgr_config = deadpool_postgres::ManagerConfig {
            recycling_method: RecyclingMethod::Fast,
        };
        let mgr = Manager::from_config(config.clone(), tls, mgr_config);
        let pool = Pool::builder(mgr).max_size(16).build().unwrap();
        #[cfg(feature = "postgresql_embedded")]
        return Ok(Self {
            db_config: db_config.clone(),
            config,
            pool,
            pg_inst,
        });
        #[cfg(not(feature = "postgresql_embedded"))]
        return Ok(Self {
            db_config: db_config.clone(),
            config,
            pool,
        });
    }

    /// Runs database migrations if needed.
    pub async fn run_migrations(&self) -> Result<(), DBError> {
        debug!("Applying database migrations if needed...");
        let mut client = self.pool.get().await?;
        let report = embedded::migrations::runner()
            .run_async(&mut **client)
            .await?;
        if report.applied_migrations().is_empty() {
            debug!(
                "Database migrations finished: {} migrations were applied",
                report.applied_migrations().len()
            );
        } else {
            info!(
                "Database migrations finished: {} migrations were applied",
                report.applied_migrations().len()
            );
        }

        // YAML -> JSON migration
        self.perform_yaml_to_json_migration().await?;

        let default_tenant = TenantRecord::default();
        self.get_or_create_tenant_id(
            default_tenant.id.0,
            default_tenant.tenant,
            default_tenant.provider,
        )
        .await?;
        Ok(())
    }

    /// Performs the conversion of the pipeline table fields which are YAML to become JSON.
    async fn perform_yaml_to_json_migration(&self) -> Result<(), DBError> {
        let mut client = self.pool.get().await?;
        let txn = client.transaction().await?;
        let stmt = txn
            .prepare_cached(
                "SELECT p.id, p.runtime_config, p.program_config, p.program_info, p.deployment_error, p.deployment_config
                 FROM pipeline AS p
                 WHERE p.uses_json = FALSE
                 ORDER BY p.id ASC"
            )
            .await?;
        let rows: Vec<Row> = txn.query(&stmt, &[]).await?;
        if !rows.is_empty() {
            let stmt_update = txn
                .prepare_cached(
                    "UPDATE pipeline
                     SET runtime_config = $1,
                         program_config = $2,
                         program_info = $3,
                         deployment_error = $4,
                         deployment_config = $5,
                         uses_json = TRUE
                     WHERE id = $6",
                )
                .await?;
            let mut rows_affected = 0;
            for row in rows {
                let pipeline_id = PipelineId(row.get(0));
                let runtime_config = serde_yaml::from_str::<RuntimeConfig>(row.get(1))
                    .expect("Deserialize RuntimeConfig from YAML");
                let program_config = serde_yaml::from_str::<ProgramConfig>(row.get(2))
                    .expect("Deserialize ProgramConfig from YAML");
                let program_info = row.get::<_, Option<String>>(3).map(|s| {
                    serde_yaml::from_str::<ProgramInfo>(&s)
                        .expect("Deserialize ProgramInfo from YAML")
                });
                let deployment_error = row.get::<_, Option<String>>(4).map(|s| {
                    serde_yaml::from_str::<ErrorResponse>(&s)
                        .expect("Deserialize ErrorResponse from YAML")
                });
                let deployment_config = row.get::<_, Option<String>>(5).map(|s| {
                    serde_yaml::from_str::<PipelineConfig>(&s)
                        .expect("Deserialize PipelineConfig from YAML")
                });
                rows_affected += txn
                    .execute(
                        &stmt_update,
                        &[
                            &serde_json::to_string(&runtime_config)
                                .expect("Serialize RuntimeConfig as JSON"),
                            &serde_json::to_string(&program_config)
                                .expect("Serialize ProgramConfig as JSON"),
                            &program_info.as_ref().map(|v| {
                                serde_json::to_string(&v).expect("Serialize ProgramInfo as JSON")
                            }),
                            &deployment_error.as_ref().map(|v| {
                                serde_json::to_string(&v).expect("Serialize ErrorResponse as JSON")
                            }),
                            &deployment_config.as_ref().map(|v| {
                                serde_json::to_string(&v).expect("Serialize PipelineConfig as JSON")
                            }),
                            &pipeline_id.0,
                        ],
                    )
                    .await?;
            }
            txn.commit().await?;
            if rows_affected > 0 {
                info!("Converted YAML to JSON for {rows_affected} row(s)");
            }
        } else {
            txn.rollback().await?;
        }
        Ok(())
    }

    /// Check if the expected DB migrations have already been run
    pub async fn check_migrations(self) -> Result<Self, DBError> {
        debug!("Checking if DB migrations have been applied");
        let mut client = self.pool.get().await?;
        let runner = embedded::migrations::runner();
        let expected_max_version = runner
            .get_migrations()
            .iter()
            .map(|m| m.version())
            .fold(0i32, |a, b| a.max(b));
        let migration = runner.get_last_applied_migration_async(&mut **client).await;
        if let Ok(Some(m)) = migration {
            let v = m.version();
            info!("Expected version = {expected_max_version}. Actual version = {v}.");
            if v == expected_max_version {
                Ok(self)
            } else {
                Err(DBError::MissingMigrations {
                    expected: expected_max_version,
                    actual: v,
                })
            }
        } else {
            info!("Expected version = {expected_max_version}. Actual version = None.");
            Err(DBError::MissingMigrations {
                expected: expected_max_version,
                actual: 0,
            })
        }
    }
}
