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

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

//! Queries for `kb_observed_tokens`.

/// Inserts or updates one observed token by mint.
pub async fn upsert_observed_token(
    database: &crate::KbDatabase,
    dto: &crate::KbObservedTokenDto,
) -> Result<i64, crate::KbError> {
    match database.connection() {
        crate::KbDatabaseConnection::Sqlite(pool) => {
            let decimals_i64 = dto.decimals.map(i64::from);
            let insert_result = sqlx::query(
                r#"
INSERT INTO kb_observed_tokens (
    mint,
    symbol,
    name,
    decimals,
    token_program,
    status,
    first_seen_at,
    last_seen_at,
    updated_at
)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(mint) DO UPDATE SET
    symbol = excluded.symbol,
    name = excluded.name,
    decimals = excluded.decimals,
    token_program = excluded.token_program,
    status = excluded.status,
    last_seen_at = excluded.last_seen_at,
    updated_at = excluded.updated_at
                "#,
            )
            .bind(dto.mint.clone())
            .bind(dto.symbol.clone())
            .bind(dto.name.clone())
            .bind(decimals_i64)
            .bind(dto.token_program.clone())
            .bind(dto.status.to_i16())
            .bind(dto.first_seen_at.to_rfc3339())
            .bind(dto.last_seen_at.to_rfc3339())
            .bind(dto.updated_at.to_rfc3339())
            .execute(pool)
            .await;
            if let Err(error) = insert_result {
                return Err(crate::KbError::Db(format!(
                    "cannot upsert kb_observed_tokens on sqlite: {}",
                    error
                )));
            }
            let select_result = sqlx::query_scalar::<sqlx::Sqlite, i64>(
                r#"
SELECT id
FROM kb_observed_tokens
WHERE mint = ?
LIMIT 1
                "#,
            )
            .bind(dto.mint.clone())
            .fetch_one(pool)
            .await;
            match select_result {
                Ok(id) => return Ok(id),
                Err(error) => {
                    return Err(crate::KbError::Db(format!(
                        "cannot fetch kb_observed_tokens id for mint '{}' on sqlite: {}",
                        dto.mint, error
                    )));
                },
            }
        },
    }
}

/// Reads one observed token by mint.
pub async fn get_observed_token_by_mint(
    database: &crate::KbDatabase,
    mint: &str,
) -> Result<std::option::Option<crate::KbObservedTokenDto>, crate::KbError> {
    match database.connection() {
        crate::KbDatabaseConnection::Sqlite(pool) => {
            let query_result = sqlx::query_as::<sqlx::Sqlite, crate::KbObservedTokenEntity>(
                r#"
SELECT
    id,
    mint,
    symbol,
    name,
    decimals,
    token_program,
    status,
    first_seen_at,
    last_seen_at,
    updated_at
FROM kb_observed_tokens
WHERE mint = ?
LIMIT 1
                    "#,
            )
            .bind(mint)
            .fetch_optional(pool)
            .await;
            let entity_option = match query_result {
                Ok(entity_option) => entity_option,
                Err(error) => {
                    return Err(crate::KbError::Db(format!(
                        "cannot read observed token '{}' on sqlite: {}",
                        mint, error
                    )));
                },
            };
            match entity_option {
                Some(entity) => {
                    let dto_result = crate::KbObservedTokenDto::try_from(entity);
                    match dto_result {
                        Ok(dto) => return Ok(Some(dto)),
                        Err(error) => return Err(error),
                    }
                },
                None => return Ok(None),
            }
        },
    }
}

/// Lists observed tokens ordered by newest first.
pub async fn list_observed_tokens(
    database: &crate::KbDatabase,
    limit: u32,
) -> Result<std::vec::Vec<crate::KbObservedTokenDto>, crate::KbError> {
    if limit == 0 {
        return Ok(std::vec::Vec::new());
    }
    match database.connection() {
        crate::KbDatabaseConnection::Sqlite(pool) => {
            let query_result = sqlx::query_as::<sqlx::Sqlite, crate::KbObservedTokenEntity>(
                r#"
SELECT
    id,
    mint,
    symbol,
    name,
    decimals,
    token_program,
    status,
    first_seen_at,
    last_seen_at,
    updated_at
FROM kb_observed_tokens
ORDER BY id DESC
LIMIT ?
                    "#,
            )
            .bind(i64::from(limit))
            .fetch_all(pool)
            .await;
            let entities = match query_result {
                Ok(entities) => entities,
                Err(error) => {
                    return Err(crate::KbError::Db(format!(
                        "cannot list observed tokens on sqlite: {}",
                        error
                    )));
                },
            };
            let mut dtos = std::vec::Vec::new();
            for entity in entities {
                let dto_result = crate::KbObservedTokenDto::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 observed_token_roundtrip_works() {
        let tempdir = tempfile::tempdir().expect("tempdir must succeed");
        let database_path = tempdir.path().join("observed_token.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 dto = crate::KbObservedTokenDto::new(
            "So11111111111111111111111111111111111111112".to_string(),
            Some("WSOL".to_string()),
            Some("Wrapped SOL".to_string()),
            Some(9),
            crate::SPL_TOKEN_PROGRAM_ID.to_string(),
            crate::KbObservedTokenStatus::Active,
        );
        let inserted_id = crate::upsert_observed_token(&database, &dto)
            .await
            .expect("upsert must succeed");
        assert!(inserted_id > 0);
        let fetched = crate::get_observed_token_by_mint(
            &database,
            "So11111111111111111111111111111111111111112",
        )
        .await
        .expect("fetch must succeed");
        assert!(fetched.is_some());
        let fetched = fetched.expect("token must exist");
        assert_eq!(fetched.symbol.as_deref(), Some("WSOL"));
        assert_eq!(fetched.decimals, Some(9));
        assert_eq!(fetched.status, crate::KbObservedTokenStatus::Active);
        let listed = crate::list_observed_tokens(&database, 10).await.expect("list must succeed");
        assert_eq!(listed.len(), 1);
    }
}