😸

SurrealDBにbytesデータを入れる

2024/12/13に公開
2

SurrealDBにbytesデータを挿入するコードの備忘録です

use std::fs::File;
use std::io::Read;
use surrealdb::Surreal;
use surrealdb::engine::local::Db;
use surrealdb::sql::Bytes;
use base64::{engine::general_purpose::STANDARD, Engine};
use tokio;
use surrealdb::engine::local::Mem;

fn csv_to_bytes(file_path: &str) -> std::io::Result<Bytes> {
    let mut file = File::open(file_path)?;
    let mut buffer = Vec::new();
    file.read_to_end(&mut buffer)?;
    Ok(Bytes::from(buffer))
}

async fn insert_bytes_into_surrealdb(db: &Surreal<Db>, bytes: Bytes) -> surrealdb::Result<()> {
    db.use_ns("example_ns").use_db("example_db").await?;

    // Base64エンコードされたデータを準備
    let encoded_data = STANDARD.encode(bytes.to_vec());

    db.query("CREATE type::thing($tb, $id) SET file_data = $data")
        .bind(("tb", "csv_data"))
        .bind(("id", "file_bytes"))
        .bind(("data", encoded_data)) // Base64文字列を保存
        .await?;
    Ok(())
}

async fn get_bytes_from_surrealdb(db: &Surreal<Db>) -> surrealdb::Result<Option<Vec<u8>>> {
    db.use_ns("example_ns").use_db("example_db").await?;

    let result: Option<serde_json::Value> = db
        .query("SELECT file_data FROM type::thing($tb, $id)")
        .bind(("tb", "csv_data"))
        .bind(("id", "file_bytes"))
        .await?
        .take(0)?;

    if let Some(value) = result {
        if let Some(encoded_data) = value.get("file_data").and_then(|v| v.as_str()) {
            let byte_data = STANDARD.decode(encoded_data).map_err(|e| e.to_string()).unwrap();
            return Ok(Some(byte_data));
        }
    }

    Ok(None)
}

#[tokio::main]
async fn main() -> surrealdb::Result<()> {
    let file_path = "/Users/ods/Documents/rust_test/sample.csv";

    // SurrealDBのインスタンスを作成
    let db = Surreal::new::<Mem>(()).await?;

    let bytes = csv_to_bytes(file_path).unwrap();

    if let Err(e) = insert_bytes_into_surrealdb(&db, bytes).await {
        eprintln!("Error inserting into SurrealDB: {:?}", e);
    }

    match get_bytes_from_surrealdb(&db).await? {
        Some(byte_data) => {
            println!("Retrieved Bytes data: {:?}", byte_data);
            match String::from_utf8(byte_data) {
                Ok(string_data) => println!("Retrieved String data: {}", string_data),
                Err(e) => eprintln!("Failed to convert Bytes to String: {}", e),
            }
        }
        None => {
            println!("No data found in SurrealDB.");
        }
    }

    Ok(())
}

Discussion