khadhroony-bobobot/kb_lib/src/db/queries/known_ws_endpoint.rs

// file: kb_lib/src/db/queries/known_ws_endpoint.rs

//! Queries for `k_sol_known_ws_endpoints`.

/// Inserts or updates one known WS endpoint row.
pub async fn query_known_ws_endpoints_upsert(
    database: &crate::Database,
    dto: &crate::KnownWsEndpointDto,
) -> Result<(), crate::Error> {
    let entity_result = crate::KnownWsEndpointEntity::try_from(dto.clone());
    let entity = match entity_result {
        Ok(entity) => entity,
        Err(error) => return Err(error),
    };
    match database.connection() {
        crate::DatabaseConnection::Sqlite(pool) => {
            let query_result = sqlx::query(
                r#"
INSERT INTO k_sol_known_ws_endpoints (
    name,
    provider,
    url,
    enabled,
    roles_json,
    last_seen_at,
    updated_at
)
VALUES (?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(name) DO UPDATE SET
    provider = excluded.provider,
    url = excluded.url,
    enabled = excluded.enabled,
    roles_json = excluded.roles_json,
    last_seen_at = excluded.last_seen_at,
    updated_at = excluded.updated_at
                "#,
            )
            .bind(entity.name)
            .bind(entity.provider)
            .bind(entity.url)
            .bind(entity.enabled)
            .bind(entity.roles_json)
            .bind(entity.last_seen_at)
            .bind(entity.updated_at)
            .execute(pool)
            .await;
            match query_result {
                Ok(_) => return Ok(()),
                Err(error) => {
                    return Err(crate::Error::Db(format!(
                        "cannot upsert k_sol_known_ws_endpoints on sqlite: {}",
                        error
                    )));
                },
            }
        },
    }
}

/// Reads one known WS endpoint by name.
pub async fn query_known_ws_endpoints_get(
    database: &crate::Database,
    name: &str,
) -> Result<std::option::Option<crate::KnownWsEndpointDto>, crate::Error> {
    match database.connection() {
        crate::DatabaseConnection::Sqlite(pool) => {
            let query_result = sqlx::query_as::<sqlx::Sqlite, crate::KnownWsEndpointEntity>(
                r#"
SELECT
    name,
    provider,
    url,
    enabled,
    roles_json,
    last_seen_at,
    updated_at
FROM k_sol_known_ws_endpoints
WHERE name = ?
LIMIT 1
                    "#,
            )
            .bind(name)
            .fetch_optional(pool)
            .await;
            let entity_option = match query_result {
                Ok(entity_option) => entity_option,
                Err(error) => {
                    return Err(crate::Error::Db(format!(
                        "cannot read known ws endpoint '{}' on sqlite: {}",
                        name, error
                    )));
                },
            };
            match entity_option {
                Some(entity) => {
                    let dto_result = crate::KnownWsEndpointDto::try_from(entity);
                    match dto_result {
                        Ok(dto) => return Ok(Some(dto)),
                        Err(error) => return Err(error),
                    }
                },
                None => return Ok(None),
            }
        },
    }
}

/// Lists all known WS endpoints.
pub async fn query_known_ws_endpoints_list(
    database: &crate::Database,
) -> Result<std::vec::Vec<crate::KnownWsEndpointDto>, crate::Error> {
    match database.connection() {
        crate::DatabaseConnection::Sqlite(pool) => {
            let query_result = sqlx::query_as::<sqlx::Sqlite, crate::KnownWsEndpointEntity>(
                r#"
SELECT
    name,
    provider,
    url,
    enabled,
    roles_json,
    last_seen_at,
    updated_at
FROM k_sol_known_ws_endpoints
ORDER BY name ASC
                    "#,
            )
            .fetch_all(pool)
            .await;
            let entities = match query_result {
                Ok(entities) => entities,
                Err(error) => {
                    return Err(crate::Error::Db(format!(
                        "cannot list known ws endpoints on sqlite: {}",
                        error
                    )));
                },
            };
            let mut dtos = std::vec::Vec::new();
            for entity in entities {
                let dto_result = crate::KnownWsEndpointDto::try_from(entity);
                let dto = match dto_result {
                    Ok(dto) => dto,
                    Err(error) => return Err(error),
                };
                dtos.push(dto);
            }
            return Ok(dtos);
        },
    }
}

#[cfg(test)]
mod tests {
    #[tokio::test]
    async fn known_ws_endpoint_roundtrip_works() {
        let tempdir = tempfile::tempdir().expect("tempdir must succeed");
        let database_path = tempdir.path().join("known_ws_endpoint.sqlite3");
        let config = crate::DatabaseConfig {
            enabled: true,
            backend: crate::DatabaseBackend::Sqlite,
            sqlite: crate::SqliteDatabaseConfig {
                path: database_path.to_string_lossy().to_string(),
                create_if_missing: true,
                busy_timeout_ms: 5000,
                max_connections: 1,
                auto_initialize_schema: true,
                use_wal: true,
            },
        };
        let database = crate::Database::connect_and_initialize(&config)
            .await
            .expect("database init must succeed");
        let dto = crate::KnownWsEndpointDto::new(
            "mainnet_public_ws_slots".to_string(),
            "solana".to_string(),
            "wss://api.mainnet.solana.com".to_string(),
            true,
            vec!["ws_slots".to_string(), "ws_subscriptions".to_string()],
        );
        crate::query_known_ws_endpoints_upsert(&database, &dto)
            .await
            .expect("upsert must succeed");
        let fetched = crate::query_known_ws_endpoints_get(&database, "mainnet_public_ws_slots")
            .await
            .expect("fetch must succeed");
        assert!(fetched.is_some());
        let fetched = fetched.expect("endpoint must exist");
        assert_eq!(fetched.provider, "solana");
        assert_eq!(fetched.roles.len(), 2);
        let listed = crate::query_known_ws_endpoints_list(&database)
            .await
            .expect("list must succeed");
        assert_eq!(listed.len(), 1);
    }
}