|
import io |
|
from collections.abc import Generator |
|
from pathlib import Path |
|
|
|
from supabase import Client |
|
|
|
from configs import dify_config |
|
from extensions.storage.base_storage import BaseStorage |
|
|
|
|
|
class SupabaseStorage(BaseStorage): |
|
"""Implementation for supabase obs storage.""" |
|
|
|
def __init__(self): |
|
super().__init__() |
|
if dify_config.SUPABASE_URL is None: |
|
raise ValueError("SUPABASE_URL is not set") |
|
if dify_config.SUPABASE_API_KEY is None: |
|
raise ValueError("SUPABASE_API_KEY is not set") |
|
if dify_config.SUPABASE_BUCKET_NAME is None: |
|
raise ValueError("SUPABASE_BUCKET_NAME is not set") |
|
|
|
self.bucket_name = dify_config.SUPABASE_BUCKET_NAME |
|
self.client = Client(supabase_url=dify_config.SUPABASE_URL, supabase_key=dify_config.SUPABASE_API_KEY) |
|
self.create_bucket(id=dify_config.SUPABASE_BUCKET_NAME, bucket_name=dify_config.SUPABASE_BUCKET_NAME) |
|
|
|
def create_bucket(self, id, bucket_name): |
|
if not self.bucket_exists(): |
|
self.client.storage.create_bucket(id=id, name=bucket_name) |
|
|
|
def save(self, filename, data): |
|
self.client.storage.from_(self.bucket_name).upload(filename, data) |
|
|
|
def load_once(self, filename: str) -> bytes: |
|
content = self.client.storage.from_(self.bucket_name).download(filename) |
|
return content |
|
|
|
def load_stream(self, filename: str) -> Generator: |
|
result = self.client.storage.from_(self.bucket_name).download(filename) |
|
byte_stream = io.BytesIO(result) |
|
while chunk := byte_stream.read(4096): |
|
yield chunk |
|
|
|
def download(self, filename, target_filepath): |
|
result = self.client.storage.from_(self.bucket_name).download(filename) |
|
Path(target_filepath).write_bytes(result) |
|
|
|
def exists(self, filename): |
|
result = self.client.storage.from_(self.bucket_name).list(filename) |
|
if result.count() > 0: |
|
return True |
|
return False |
|
|
|
def delete(self, filename): |
|
self.client.storage.from_(self.bucket_name).remove(filename) |
|
|
|
def bucket_exists(self): |
|
buckets = self.client.storage.list_buckets() |
|
return any(bucket.name == self.bucket_name for bucket in buckets) |
|
|
