headless_lms_chatbot/
azure_blob_storage.rs

1use std::collections::HashMap;
2
3use crate::prelude::*;
4use anyhow::Context;
5use azure_core::prelude::Metadata;
6use azure_storage::StorageCredentials;
7use azure_storage_blobs::prelude::*;
8use bytes::Bytes;
9use futures::StreamExt;
10use headless_lms_utils::{ApplicationConfiguration, AzureBlobStorageConfiguration};
11
12/// A client for interacting with Azure Blob Storage.
13pub struct AzureBlobClient {
14    container_client: ContainerClient,
15    pub container_name: String,
16}
17
18impl AzureBlobClient {
19    pub async fn new(
20        app_config: &ApplicationConfiguration,
21        container_name: &str,
22    ) -> anyhow::Result<Self> {
23        let azure_configuration = app_config
24            .azure_configuration
25            .as_ref()
26            .context("Azure configuration is missing")?;
27        let AzureBlobStorageConfiguration {
28            storage_account,
29            access_key,
30        } = azure_configuration
31            .blob_storage_config
32            .clone()
33            .context("Azure Blob Storage configuration is missing")?;
34
35        let container_name = container_name.to_string();
36
37        let storage_credentials = StorageCredentials::access_key(&storage_account, access_key);
38        let blob_service_client = BlobServiceClient::new(storage_account, storage_credentials);
39        let container_client = blob_service_client.container_client(container_name.clone());
40
41        Ok(AzureBlobClient {
42            container_client,
43            container_name,
44        })
45    }
46
47    /// Ensures the container used to store the blobs exists. If it does not, the container is created.
48    pub async fn ensure_container_exists(&self) -> anyhow::Result<()> {
49        if self.container_client.exists().await? {
50            return Ok(());
51        }
52
53        info!(
54            "Azure blob storage container '{}' does not exist. Creating...",
55            self.container_client.container_name()
56        );
57        self.container_client
58            .create()
59            .public_access(PublicAccess::None)
60            .await?;
61        Ok(())
62    }
63
64    /// Uploads a file to the specified container.
65    pub async fn upload_file(
66        &self,
67        blob_path: &str,
68        file_bytes: &[u8],
69        metadata: Option<HashMap<String, Bytes>>,
70    ) -> anyhow::Result<()> {
71        let blob_client = self.container_client.blob_client(blob_path);
72
73        let mut put_blob = blob_client.put_block_blob(file_bytes.to_vec());
74
75        if let Some(meta) = metadata {
76            let mut m = Metadata::new();
77            for (key, value) in meta {
78                m.insert(key, value);
79            }
80            put_blob = put_blob.metadata(m);
81        }
82
83        put_blob.await?;
84
85        info!("Blob '{}' uploaded successfully.", blob_path);
86        Ok(())
87    }
88
89    /// Deletes a file (blob) from the specified container.
90    pub async fn delete_file(&self, path: &str) -> anyhow::Result<()> {
91        let blob_client = self.container_client.blob_client(path);
92
93        blob_client.delete().await?;
94
95        info!("Blob '{}' deleted successfully.", path);
96        Ok(())
97    }
98
99    pub async fn list_files_with_prefix(&self, prefix: &str) -> anyhow::Result<Vec<String>> {
100        let mut result = Vec::new();
101        let prefix_owned = prefix.to_string();
102        let response = self.container_client.list_blobs().prefix(prefix_owned);
103        let mut stream = response.into_stream();
104
105        while let Some(list) = stream.next().await {
106            let list = list?;
107            let blobs: Vec<_> = list.blobs.blobs().collect();
108            for blob in blobs {
109                result.push(blob.name.clone());
110            }
111        }
112        Ok(result)
113    }
114}