Austin/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2024-11-16 11:42:29 +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 collections.abc import Generator
from contextlib import closing
import oss2 as aliyun_s3 import oss2 as aliyun_s3
from flask import Flask from flask import Flask
@ -34,15 +33,15 @@ class AliyunOssStorage(BaseStorage):
self.client.put_object(self.__wrapper_folder_filename(filename), data) self.client.put_object(self.__wrapper_folder_filename(filename), data)
def load_once(self, filename: str) -> bytes: def load_once(self, filename: str) -> bytes:
with closing(self.client.get_object(self.__wrapper_folder_filename(filename))) as obj: obj = self.client.get_object(self.__wrapper_folder_filename(filename))
data = obj.read() data = obj.read()
return data return data
def load_stream(self, filename: str) -> Generator: def load_stream(self, filename: str) -> Generator:
def generate(filename: str = filename) -> Generator: def generate(filename: str = filename) -> Generator:
with closing(self.client.get_object(self.__wrapper_folder_filename(filename))) as obj: obj = self.client.get_object(self.__wrapper_folder_filename(filename))
while chunk := obj.read(4096): while chunk := obj.read(4096):
yield chunk yield chunk
return generate() return generate()

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

@ -1,6 +1,5 @@
import logging import logging
from collections.abc import Generator from collections.abc import Generator
from contextlib import closing
import boto3 import boto3
from botocore.client import Config from botocore.client import Config
@ -55,8 +54,7 @@ class AwsS3Storage(BaseStorage):
def load_once(self, filename: str) -> bytes: def load_once(self, filename: str) -> bytes:
try: try:
with closing(self.client) as client: data = self.client.get_object(Bucket=self.bucket_name, Key=filename)["Body"].read()
data = client.get_object(Bucket=self.bucket_name, Key=filename)["Body"].read()
except ClientError as ex: except ClientError as ex:
if ex.response["Error"]["Code"] == "NoSuchKey": if ex.response["Error"]["Code"] == "NoSuchKey":
raise FileNotFoundError("File not found") raise FileNotFoundError("File not found")
@ -67,9 +65,8 @@ class AwsS3Storage(BaseStorage):
def load_stream(self, filename: str) -> Generator: def load_stream(self, filename: str) -> Generator:
def generate(filename: str = filename) -> Generator: def generate(filename: str = filename) -> Generator:
try: try:
with closing(self.client) as client: response = self.client.get_object(Bucket=self.bucket_name, Key=filename)
response = client.get_object(Bucket=self.bucket_name, Key=filename) yield from response["Body"].iter_chunks()
yield from response["Body"].iter_chunks()
except ClientError as ex: except ClientError as ex:
if ex.response["Error"]["Code"] == "NoSuchKey": if ex.response["Error"]["Code"] == "NoSuchKey":
raise FileNotFoundError("File not found") raise FileNotFoundError("File not found")
@ -79,16 +76,14 @@ class AwsS3Storage(BaseStorage):
return generate() return generate()
def download(self, filename, target_filepath): def download(self, filename, target_filepath):
with closing(self.client) as client: self.client.download_file(self.bucket_name, filename, target_filepath)
client.download_file(self.bucket_name, filename, target_filepath)
def exists(self, filename): def exists(self, filename):
with closing(self.client) as client: try:
try: self.client.head_object(Bucket=self.bucket_name, Key=filename)
client.head_object(Bucket=self.bucket_name, Key=filename) return True
return True except:
except: return False
return False
def delete(self, filename): def delete(self, filename):
self.client.delete_object(Bucket=self.bucket_name, Key=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 io
import json import json
from collections.abc import Generator from collections.abc import Generator
from contextlib import closing
from flask import Flask from flask import Flask
from google.cloud import storage as google_cloud_storage from google.cloud import storage as google_cloud_storage
@ -43,7 +42,7 @@ class GoogleCloudStorage(BaseStorage):
def generate(filename: str = filename) -> Generator: def generate(filename: str = filename) -> Generator:
bucket = self.client.get_bucket(self.bucket_name) bucket = self.client.get_bucket(self.bucket_name)
blob = bucket.get_blob(filename) 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): while chunk := blob_stream.read(4096):
yield chunk yield chunk

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

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