// file: kb_lib/src/db/queries/pool_token.rs //! Queries for `kb_pool_tokens`. /// Inserts or updates one normalized pool token composition row. pub async fn upsert_pool_token( database: &crate::KbDatabase, dto: &crate::KbPoolTokenDto, ) -> Result { match database.connection() { crate::KbDatabaseConnection::Sqlite(pool) => { let query_result = sqlx::query( r#" INSERT INTO kb_pool_tokens ( pool_id, token_id, role, vault_address, token_order, created_at, updated_at ) VALUES (?, ?, ?, ?, ?, ?, ?) ON CONFLICT(pool_id, token_id, role) DO UPDATE SET vault_address = excluded.vault_address, token_order = excluded.token_order, updated_at = excluded.updated_at "#, ) .bind(dto.pool_id) .bind(dto.token_id) .bind(dto.role.to_i16()) .bind(dto.vault_address.clone()) .bind(dto.token_order) .bind(dto.created_at.to_rfc3339()) .bind(dto.updated_at.to_rfc3339()) .execute(pool) .await; if let Err(error) = query_result { return Err(crate::KbError::Db(format!( "cannot upsert kb_pool_tokens on sqlite: {}", error ))); } let id_result = sqlx::query_scalar::( r#" SELECT id FROM kb_pool_tokens WHERE pool_id = ? AND token_id = ? AND role = ? LIMIT 1 "#, ) .bind(dto.pool_id) .bind(dto.token_id) .bind(dto.role.to_i16()) .fetch_one(pool) .await; match id_result { Ok(id) => return Ok(id), Err(error) => { return Err(crate::KbError::Db(format!( "cannot fetch kb_pool_tokens id on sqlite: {}", error ))); }, } }, } } /// Lists normalized pool token rows by pool id. pub async fn list_pool_tokens_by_pool_id( database: &crate::KbDatabase, pool_id: i64, ) -> Result, crate::KbError> { match database.connection() { crate::KbDatabaseConnection::Sqlite(pool) => { let query_result = sqlx::query_as::( r#" SELECT id, pool_id, token_id, role, vault_address, token_order, created_at, updated_at FROM kb_pool_tokens WHERE pool_id = ? ORDER BY token_order ASC, id ASC "#, ) .bind(pool_id) .fetch_all(pool) .await; let entities = match query_result { Ok(entities) => entities, Err(error) => { return Err(crate::KbError::Db(format!( "cannot list kb_pool_tokens for pool_id '{}' on sqlite: {}", pool_id, error ))); }, }; let mut dtos = std::vec::Vec::new(); for entity in entities { let dto_result = crate::KbPoolTokenDto::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 pool_token_roundtrip_works() { let tempdir = tempfile::tempdir().expect("tempdir must succeed"); let database_path = tempdir.path().join("pool_token_roundtrip.sqlite3"); let config = crate::KbDatabaseConfig { enabled: true, backend: crate::KbDatabaseBackend::Sqlite, sqlite: crate::KbSqliteDatabaseConfig { 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::KbDatabase::connect_and_initialize(&config) .await .expect("database init must succeed"); let dex_id = crate::upsert_dex( &database, &crate::KbDexDto::new("raydium".to_string(), "Raydium".to_string(), None, None, true), ) .await .expect("dex upsert must succeed"); let token_id = crate::upsert_token( &database, &crate::KbTokenDto::new( "Base111111111111111111111111111111111111111".to_string(), Some("BASE".to_string()), Some("Base Token".to_string()), Some(6), crate::SPL_TOKEN_PROGRAM_ID.to_string(), false, ), ) .await .expect("token upsert must succeed"); let pool_id = crate::upsert_pool( &database, &crate::KbPoolDto::new( dex_id, "Pool111111111111111111111111111111111111111".to_string(), crate::KbPoolKind::Amm, crate::KbPoolStatus::Active, ), ) .await .expect("pool upsert must succeed"); let pool_token_id = crate::upsert_pool_token( &database, &crate::KbPoolTokenDto::new( pool_id, token_id, crate::KbPoolTokenRole::Base, Some("Vault111111111111111111111111111111111111".to_string()), Some(0), ), ) .await .expect("pool token upsert must succeed"); assert!(pool_token_id > 0); let pool_tokens = crate::list_pool_tokens_by_pool_id(&database, pool_id) .await .expect("list pool tokens must succeed"); assert_eq!(pool_tokens.len(), 1); assert_eq!(pool_tokens[0].role, crate::KbPoolTokenRole::Base); } }