Underground-Digital/Workflow-Engine

1import io2from collections.abc import Generator3from pathlib import Path4 5from supabase import Client6 7from configs import dify_config8from extensions.storage.base_storage import BaseStorage9 10 11class SupabaseStorage(BaseStorage):12    """Implementation for supabase obs storage."""13 14    def __init__(self):15        super().__init__()16        if dify_config.SUPABASE_URL is None:17            raise ValueError("SUPABASE_URL is not set")18        if dify_config.SUPABASE_API_KEY is None:19            raise ValueError("SUPABASE_API_KEY is not set")20        if dify_config.SUPABASE_BUCKET_NAME is None:21            raise ValueError("SUPABASE_BUCKET_NAME is not set")22 23        self.bucket_name = dify_config.SUPABASE_BUCKET_NAME24        self.client = Client(supabase_url=dify_config.SUPABASE_URL, supabase_key=dify_config.SUPABASE_API_KEY)25        self.create_bucket(id=dify_config.SUPABASE_BUCKET_NAME, bucket_name=dify_config.SUPABASE_BUCKET_NAME)26 27    def create_bucket(self, id, bucket_name):28        if not self.bucket_exists():29            self.client.storage.create_bucket(id=id, name=bucket_name)30 31    def save(self, filename, data):32        self.client.storage.from_(self.bucket_name).upload(filename, data)33 34    def load_once(self, filename: str) -> bytes:35        content = self.client.storage.from_(self.bucket_name).download(filename)36        return content37 38    def load_stream(self, filename: str) -> Generator:39        result = self.client.storage.from_(self.bucket_name).download(filename)40        byte_stream = io.BytesIO(result)41        while chunk := byte_stream.read(4096):  # Read in chunks of 4KB42            yield chunk43 44    def download(self, filename, target_filepath):45        result = self.client.storage.from_(self.bucket_name).download(filename)46        Path(target_filepath).write_bytes(result)47 48    def exists(self, filename):49        result = self.client.storage.from_(self.bucket_name).list(filename)50        if result.count() > 0:51            return True52        return False53 54    def delete(self, filename):55        self.client.storage.from_(self.bucket_name).remove(filename)56 57    def bucket_exists(self):58        buckets = self.client.storage.list_buckets()59        return any(bucket.name == self.bucket_name for bucket in buckets)60