test/dify
1
0
Fork 0
mirror of https://gitee.com/dify_ai/dify.git synced 2024-12-03 11:48:08 +08:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

refactor: remove unnecessary 'closing' usage for boto3 client (#9343)

Browse Source
This commit is contained in:
omr 2024-10-15 09:42:39 +09:00 committed by GitHub
parent 5eb00502ec
commit 6d2c6caa23
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 24 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
api/extensions/storage/aliyun_oss_storage.py
View File

@ -1,5 +1,4 @@
from collections.abc import Generator
from contextlib import closing
import oss2 as aliyun_s3
from flask import Flask
@ -34,15 +33,15 @@ class AliyunOssStorage(BaseStorage):
self.client.put_object(self.__wrapper_folder_filename(filename), data)
def load_once(self, filename: str) -> bytes:
with closing(self.client.get_object(self.__wrapper_folder_filename(filename))) as obj:
data = obj.read()
obj = self.client.get_object(self.__wrapper_folder_filename(filename))
data = obj.read()
return data
def load_stream(self, filename: str) -> Generator:
def generate(filename: str = filename) -> Generator:
with closing(self.client.get_object(self.__wrapper_folder_filename(filename))) as obj:
while chunk := obj.read(4096):
yield chunk
obj = self.client.get_object(self.__wrapper_folder_filename(filename))
while chunk := obj.read(4096):
yield chunk
return generate()

23
api/extensions/storage/aws_s3_storage.py
View File

@ -1,6 +1,5 @@
import logging
from collections.abc import Generator
from contextlib import closing
import boto3
from botocore.client import Config
@ -55,8 +54,7 @@ class AwsS3Storage(BaseStorage):
def load_once(self, filename: str) -> bytes:
try:
with closing(self.client) as client:
data = client.get_object(Bucket=self.bucket_name, Key=filename)["Body"].read()
data = self.client.get_object(Bucket=self.bucket_name, Key=filename)["Body"].read()
except ClientError as ex:
if ex.response["Error"]["Code"] == "NoSuchKey":
raise FileNotFoundError("File not found")
@ -67,9 +65,8 @@ class AwsS3Storage(BaseStorage):
def load_stream(self, filename: str) -> Generator:
def generate(filename: str = filename) -> Generator:
try:
with closing(self.client) as client:
response = client.get_object(Bucket=self.bucket_name, Key=filename)
yield from response["Body"].iter_chunks()
response = self.client.get_object(Bucket=self.bucket_name, Key=filename)
yield from response["Body"].iter_chunks()
except ClientError as ex:
if ex.response["Error"]["Code"] == "NoSuchKey":
raise FileNotFoundError("File not found")
@ -79,16 +76,14 @@ class AwsS3Storage(BaseStorage):
return generate()
def download(self, filename, target_filepath):
with closing(self.client) as client:
client.download_file(self.bucket_name, filename, target_filepath)
self.client.download_file(self.bucket_name, filename, target_filepath)
def exists(self, filename):
with closing(self.client) as client:
try:
client.head_object(Bucket=self.bucket_name, Key=filename)
return True
except:
return False
try:
self.client.head_object(Bucket=self.bucket_name, Key=filename)
return True
except:
return False
def delete(self, filename):
self.client.delete_object(Bucket=self.bucket_name, Key=filename)

3
api/extensions/storage/google_cloud_storage.py
View File

@ -2,7 +2,6 @@ import base64
import io
import json
from collections.abc import Generator
from contextlib import closing
from flask import Flask
from google.cloud import storage as google_cloud_storage
@ -43,7 +42,7 @@ class GoogleCloudStorage(BaseStorage):
def generate(filename: str = filename) -> Generator:
bucket = self.client.get_bucket(self.bucket_name)
blob = bucket.get_blob(filename)
with closing(blob.open(mode="rb")) as blob_stream:
with blob.open(mode="rb") as blob_stream:
while chunk := blob_stream.read(4096):
yield chunk

23
api/extensions/storage/oracle_oci_storage.py
View File

@ -1,5 +1,4 @@
from collections.abc import Generator
from contextlib import closing
import boto3
from botocore.exceptions import ClientError
@ -28,8 +27,7 @@ class OracleOCIStorage(BaseStorage):
def load_once(self, filename: str) -> bytes:
try:
with closing(self.client) as client:
data = client.get_object(Bucket=self.bucket_name, Key=filename)["Body"].read()
data = self.client.get_object(Bucket=self.bucket_name, Key=filename)["Body"].read()
except ClientError as ex:
if ex.response["Error"]["Code"] == "NoSuchKey":
raise FileNotFoundError("File not found")
@ -40,9 +38,8 @@ class OracleOCIStorage(BaseStorage):
def load_stream(self, filename: str) -> Generator:
def generate(filename: str = filename) -> Generator:
try:
with closing(self.client) as client:
response = client.get_object(Bucket=self.bucket_name, Key=filename)
yield from response["Body"].iter_chunks()
response = self.client.get_object(Bucket=self.bucket_name, Key=filename)
yield from response["Body"].iter_chunks()
except ClientError as ex:
if ex.response["Error"]["Code"] == "NoSuchKey":
raise FileNotFoundError("File not found")
@ -52,16 +49,14 @@ class OracleOCIStorage(BaseStorage):
return generate()
def download(self, filename, target_filepath):
with closing(self.client) as client:
client.download_file(self.bucket_name, filename, target_filepath)
self.client.download_file(self.bucket_name, filename, target_filepath)
def exists(self, filename):
with closing(self.client) as client:
try:
client.head_object(Bucket=self.bucket_name, Key=filename)
return True
except:
return False
try:
self.client.head_object(Bucket=self.bucket_name, Key=filename)
return True
except:
return False
def delete(self, filename):
self.client.delete_object(Bucket=self.bucket_name, Key=filename)