File size: 2,156 Bytes
a8b3f00 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
import os
import shutil
from collections.abc import Generator
from pathlib import Path
from flask import current_app
from configs import dify_config
from extensions.storage.base_storage import BaseStorage
class LocalFsStorage(BaseStorage):
"""Implementation for local filesystem storage."""
def __init__(self):
super().__init__()
folder = dify_config.STORAGE_LOCAL_PATH
if not os.path.isabs(folder):
folder = os.path.join(current_app.root_path, folder)
self.folder = folder
def _build_filepath(self, filename: str) -> str:
"""Build the full file path based on the folder and filename."""
if not self.folder or self.folder.endswith("/"):
return self.folder + filename
else:
return self.folder + "/" + filename
def save(self, filename, data):
filepath = self._build_filepath(filename)
folder = os.path.dirname(filepath)
os.makedirs(folder, exist_ok=True)
Path(os.path.join(os.getcwd(), filepath)).write_bytes(data)
def load_once(self, filename: str) -> bytes:
filepath = self._build_filepath(filename)
if not os.path.exists(filepath):
raise FileNotFoundError("File not found")
return Path(filepath).read_bytes()
def load_stream(self, filename: str) -> Generator:
filepath = self._build_filepath(filename)
if not os.path.exists(filepath):
raise FileNotFoundError("File not found")
with open(filepath, "rb") as f:
while chunk := f.read(4096): # Read in chunks of 4KB
yield chunk
def download(self, filename, target_filepath):
filepath = self._build_filepath(filename)
if not os.path.exists(filepath):
raise FileNotFoundError("File not found")
shutil.copyfile(filepath, target_filepath)
def exists(self, filename):
filepath = self._build_filepath(filename)
return os.path.exists(filepath)
def delete(self, filename):
filepath = self._build_filepath(filename)
if os.path.exists(filepath):
os.remove(filepath)
|