tcvectordb_rust/

client.rs

use reqwest::{Client, Response};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::time::Duration;
use url::Url;

use crate::database::Database;
use crate::enums::ReadConsistency;
use crate::error::{Result, VectorDBError};

#[derive(Debug, Clone)]
pub struct VectorDBClient {
    client: Client,
    base_url: Url,
    username: String,
    key: String,
    read_consistency: ReadConsistency,
}

#[derive(Debug, Serialize)]
#[allow(dead_code)]
struct AuthRequest {
    username: String,
    key: String,
}

#[derive(Debug, Deserialize)]
#[allow(dead_code)]
struct AuthResponse {
    code: i32,
    msg: String,
    #[serde(default)]
    token: Option<String>,
}

#[derive(Debug, Deserialize)]
struct ApiResponse<T> {
    code: i32,
    msg: String,
    #[serde(default)]
    data: Option<T>,
}

impl VectorDBClient {
    /// Create a new VectorDB client
    pub fn new(
        url: impl AsRef<str>,
        username: impl Into<String>,
        key: impl Into<String>,
        read_consistency: ReadConsistency,
        timeout_secs: u64,
    ) -> Result<Self> {
        let base_url = Url::parse(url.as_ref())?;
        let client = Client::builder()
            .timeout(Duration::from_secs(timeout_secs))
            .build()?;

        Ok(Self {
            client,
            base_url,
            username: username.into(),
            key: key.into(),
            read_consistency,
        })
    }

    /// Create a new VectorDB client with custom reqwest client
    pub fn with_client(
        client: Client,
        url: impl AsRef<str>,
        username: impl Into<String>,
        key: impl Into<String>,
        read_consistency: ReadConsistency,
    ) -> Result<Self> {
        let base_url = Url::parse(url.as_ref())?;

        Ok(Self {
            client,
            base_url,
            username: username.into(),
            key: key.into(),
            read_consistency,
        })
    }

    /// Get the HTTP client
    pub fn http_client(&self) -> &Client {
        &self.client
    }

    /// Get the base URL
    pub fn base_url(&self) -> &Url {
        &self.base_url
    }

    /// Get the username
    pub fn username(&self) -> &str {
        &self.username
    }

    /// Get the API key
    pub fn key(&self) -> &str {
        &self.key
    }

    /// Get the read consistency setting
    pub fn read_consistency(&self) -> &ReadConsistency {
        &self.read_consistency
    }

    /// Check if a database exists
    pub async fn exists_database(&self, database_name: &str) -> Result<bool> {
        let databases = self.list_databases().await?;
        Ok(databases.iter().any(|db| db.name() == database_name))
    }

    /// Create a new database
    pub async fn create_database(&self, database_name: impl Into<String>) -> Result<Database> {
        let database_name = database_name.into();
        let db = Database::new(self.clone(), database_name.clone());
        db.create().await?;
        Ok(db)
    }

    /// Create a database if it doesn't exist
    pub async fn create_database_if_not_exists(&self, database_name: impl Into<String>) -> Result<Database> {
        let database_name = database_name.into();
        if self.exists_database(&database_name).await? {
            Ok(Database::new(self.clone(), database_name))
        } else {
            self.create_database(database_name).await
        }
    }

    /// Drop a database
    pub async fn drop_database(&self, database_name: impl Into<String>) -> Result<Value> {
        let database_name = database_name.into();
        let db = Database::new(self.clone(), database_name);
        db.drop().await
    }

    /// List all databases
    pub async fn list_databases(&self) -> Result<Vec<Database>> {
        let url = self.base_url.join("/database/list")?;
        let response = self.client
            .get(url)
            .header("Authorization", format!("Bearer {}", self.key))
            .send()
            .await?;

        let api_response: ApiResponse<Value> = self.handle_response(response).await?;
        
        if api_response.code != 0 {
            return Err(VectorDBError::server_error(api_response.code, api_response.msg));
        }

        let databases_data = api_response.data.unwrap_or(Value::Array(vec![]));
        let database_names: Vec<String> = serde_json::from_value(databases_data)?;
        
        Ok(database_names
            .into_iter()
            .map(|name| Database::new(self.clone(), name))
            .collect())
    }

    /// Get a database by name
    pub async fn database(&self, database_name: impl Into<String>) -> Result<Database> {
        let database_name = database_name.into();
        if self.exists_database(&database_name).await? {
            Ok(Database::new(self.clone(), database_name))
        } else {
            Err(VectorDBError::param_error(
                14100,
                format!("Database not exist: {}", database_name),
            ))
        }
    }

    /// Internal method to handle HTTP responses
    pub(crate) async fn handle_response<T>(&self, response: Response) -> Result<T>
    where
        T: for<'de> Deserialize<'de>,
    {
        let status = response.status();
        let text = response.text().await?;

        if !status.is_success() {
            return Err(VectorDBError::connect_error(
                status.as_u16() as i32,
                format!("HTTP error {}: {}", status, text),
            ));
        }

        serde_json::from_str(&text).map_err(|e| {
            VectorDBError::unexpected_error(format!("Failed to parse response: {}", e))
        })
    }

    /// Internal method to make authenticated requests
    pub(crate) async fn request(&self, method: reqwest::Method, path: &str) -> Result<reqwest::RequestBuilder> {
        let url = self.base_url.join(path)?;
        Ok(self.client
            .request(method, url)
            .header("Authorization", format!("Bearer {}", self.key))
            .header("Content-Type", "application/json"))
    }

    /// Internal method to make GET requests
    pub(crate) async fn get(&self, path: &str) -> Result<reqwest::RequestBuilder> {
        self.request(reqwest::Method::GET, path).await
    }

    /// Internal method to make POST requests
    pub(crate) async fn post(&self, path: &str) -> Result<reqwest::RequestBuilder> {
        self.request(reqwest::Method::POST, path).await
    }

    /// Internal method to make PUT requests
    #[allow(dead_code)]
    pub(crate) async fn put(&self, path: &str) -> Result<reqwest::RequestBuilder> {
        self.request(reqwest::Method::PUT, path).await
    }

    /// Internal method to make DELETE requests
    pub(crate) async fn delete(&self, path: &str) -> Result<reqwest::RequestBuilder> {
        self.request(reqwest::Method::DELETE, path).await
    }
}