将flask改成fastapi
This commit is contained in:
208
rag/utils/s3_conn.py
Normal file
208
rag/utils/s3_conn.py
Normal file
@@ -0,0 +1,208 @@
|
||||
#
|
||||
# Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
import logging
|
||||
import boto3
|
||||
from botocore.exceptions import ClientError
|
||||
from botocore.config import Config
|
||||
import time
|
||||
from io import BytesIO
|
||||
from rag.utils import singleton
|
||||
from rag import settings
|
||||
|
||||
@singleton
|
||||
class RAGFlowS3:
|
||||
def __init__(self):
|
||||
self.conn = None
|
||||
self.s3_config = settings.S3
|
||||
self.access_key = self.s3_config.get('access_key', None)
|
||||
self.secret_key = self.s3_config.get('secret_key', None)
|
||||
self.session_token = self.s3_config.get('session_token', None)
|
||||
self.region_name = self.s3_config.get('region_name', None)
|
||||
self.endpoint_url = self.s3_config.get('endpoint_url', None)
|
||||
self.signature_version = self.s3_config.get('signature_version', None)
|
||||
self.addressing_style = self.s3_config.get('addressing_style', None)
|
||||
self.bucket = self.s3_config.get('bucket', None)
|
||||
self.prefix_path = self.s3_config.get('prefix_path', None)
|
||||
self.__open__()
|
||||
|
||||
@staticmethod
|
||||
def use_default_bucket(method):
|
||||
def wrapper(self, bucket, *args, **kwargs):
|
||||
# If there is a default bucket, use the default bucket
|
||||
actual_bucket = self.bucket if self.bucket else bucket
|
||||
return method(self, actual_bucket, *args, **kwargs)
|
||||
return wrapper
|
||||
|
||||
@staticmethod
|
||||
def use_prefix_path(method):
|
||||
def wrapper(self, bucket, fnm, *args, **kwargs):
|
||||
# If the prefix path is set, use the prefix path.
|
||||
# The bucket passed from the upstream call is
|
||||
# used as the file prefix. This is especially useful when you're using the default bucket
|
||||
if self.prefix_path:
|
||||
fnm = f"{self.prefix_path}/{bucket}/{fnm}"
|
||||
return method(self, bucket, fnm, *args, **kwargs)
|
||||
return wrapper
|
||||
|
||||
def __open__(self):
|
||||
try:
|
||||
if self.conn:
|
||||
self.__close__()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
try:
|
||||
s3_params = {}
|
||||
config_kwargs = {}
|
||||
# if not set ak/sk, boto3 s3 client would try several ways to do the authentication
|
||||
# see doc: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html#configuring-credentials
|
||||
if self.access_key and self.secret_key:
|
||||
s3_params = {
|
||||
'aws_access_key_id': self.access_key,
|
||||
'aws_secret_access_key': self.secret_key,
|
||||
'aws_session_token': self.session_token,
|
||||
}
|
||||
if self.region_name:
|
||||
s3_params['region_name'] = self.region_name
|
||||
if self.endpoint_url:
|
||||
s3_params['endpoint_url'] = self.endpoint_url
|
||||
|
||||
# Configure signature_version and addressing_style through Config object
|
||||
if self.signature_version:
|
||||
config_kwargs['signature_version'] = self.signature_version
|
||||
if self.addressing_style:
|
||||
config_kwargs['s3'] = {'addressing_style': self.addressing_style}
|
||||
|
||||
if config_kwargs:
|
||||
s3_params['config'] = Config(**config_kwargs)
|
||||
|
||||
self.conn = [boto3.client('s3', **s3_params)]
|
||||
except Exception:
|
||||
logging.exception(f"Fail to connect at region {self.region_name} or endpoint {self.endpoint_url}")
|
||||
|
||||
def __close__(self):
|
||||
del self.conn[0]
|
||||
self.conn = None
|
||||
|
||||
@use_default_bucket
|
||||
def bucket_exists(self, bucket, *args, **kwargs):
|
||||
try:
|
||||
logging.debug(f"head_bucket bucketname {bucket}")
|
||||
self.conn[0].head_bucket(Bucket=bucket)
|
||||
exists = True
|
||||
except ClientError:
|
||||
logging.exception(f"head_bucket error {bucket}")
|
||||
exists = False
|
||||
return exists
|
||||
|
||||
def health(self):
|
||||
bucket = self.bucket
|
||||
fnm = "txtxtxtxt1"
|
||||
fnm, binary = f"{self.prefix_path}/{fnm}" if self.prefix_path else fnm, b"_t@@@1"
|
||||
if not self.bucket_exists(bucket):
|
||||
self.conn[0].create_bucket(Bucket=bucket)
|
||||
logging.debug(f"create bucket {bucket} ********")
|
||||
|
||||
r = self.conn[0].upload_fileobj(BytesIO(binary), bucket, fnm)
|
||||
return r
|
||||
|
||||
def get_properties(self, bucket, key):
|
||||
return {}
|
||||
|
||||
def list(self, bucket, dir, recursive=True):
|
||||
return []
|
||||
|
||||
@use_prefix_path
|
||||
@use_default_bucket
|
||||
def put(self, bucket, fnm, binary, *args, **kwargs):
|
||||
logging.debug(f"bucket name {bucket}; filename :{fnm}:")
|
||||
for _ in range(1):
|
||||
try:
|
||||
if not self.bucket_exists(bucket):
|
||||
self.conn[0].create_bucket(Bucket=bucket)
|
||||
logging.info(f"create bucket {bucket} ********")
|
||||
r = self.conn[0].upload_fileobj(BytesIO(binary), bucket, fnm)
|
||||
|
||||
return r
|
||||
except Exception:
|
||||
logging.exception(f"Fail put {bucket}/{fnm}")
|
||||
self.__open__()
|
||||
time.sleep(1)
|
||||
|
||||
@use_prefix_path
|
||||
@use_default_bucket
|
||||
def rm(self, bucket, fnm, *args, **kwargs):
|
||||
try:
|
||||
self.conn[0].delete_object(Bucket=bucket, Key=fnm)
|
||||
except Exception:
|
||||
logging.exception(f"Fail rm {bucket}/{fnm}")
|
||||
|
||||
@use_prefix_path
|
||||
@use_default_bucket
|
||||
def get(self, bucket, fnm, *args, **kwargs):
|
||||
for _ in range(1):
|
||||
try:
|
||||
r = self.conn[0].get_object(Bucket=bucket, Key=fnm)
|
||||
object_data = r['Body'].read()
|
||||
return object_data
|
||||
except Exception:
|
||||
logging.exception(f"fail get {bucket}/{fnm}")
|
||||
self.__open__()
|
||||
time.sleep(1)
|
||||
return
|
||||
|
||||
@use_prefix_path
|
||||
@use_default_bucket
|
||||
def obj_exist(self, bucket, fnm, *args, **kwargs):
|
||||
try:
|
||||
if self.conn[0].head_object(Bucket=bucket, Key=fnm):
|
||||
return True
|
||||
except ClientError as e:
|
||||
if e.response['Error']['Code'] == '404':
|
||||
return False
|
||||
else:
|
||||
raise
|
||||
|
||||
@use_prefix_path
|
||||
@use_default_bucket
|
||||
def get_presigned_url(self, bucket, fnm, expires, *args, **kwargs):
|
||||
for _ in range(10):
|
||||
try:
|
||||
r = self.conn[0].generate_presigned_url('get_object',
|
||||
Params={'Bucket': bucket,
|
||||
'Key': fnm},
|
||||
ExpiresIn=expires)
|
||||
|
||||
return r
|
||||
except Exception:
|
||||
logging.exception(f"fail get url {bucket}/{fnm}")
|
||||
self.__open__()
|
||||
time.sleep(1)
|
||||
return
|
||||
|
||||
@use_default_bucket
|
||||
def rm_bucket(self, bucket, *args, **kwargs):
|
||||
for conn in self.conn:
|
||||
try:
|
||||
if not conn.bucket_exists(bucket):
|
||||
continue
|
||||
for o in conn.list_objects_v2(Bucket=bucket):
|
||||
conn.delete_object(bucket, o.object_name)
|
||||
conn.delete_bucket(Bucket=bucket)
|
||||
return
|
||||
except Exception as e:
|
||||
logging.error(f"Fail rm {bucket}: " + str(e))
|
||||
Reference in New Issue
Block a user