24024
af7c11d7f9
- 新增图像生成接口，支持试用、积分和自定义API Key模式 - 实现生成图片结果异步上传至MinIO存储，带重试机制 - 优化积分预扣除和异常退还逻辑，保障用户积分准确 - 添加获取生成历史记录接口，支持时间范围和分页 - 提供本地字典配置接口，支持模型、比例、提示模板和尺寸 - 实现图片批量上传接口，支持S3兼容对象存储 feat(admin): 增加管理员角色管理与权限分配接口 - 实现角色列表查询、角色创建、更新及删除功能 - 增加权限列表查询接口 - 实现用户角色分配接口，便于统一管理用户权限 - 增加系统字典增删查改接口，支持分类过滤和排序 - 权限控制全面覆盖管理接口，保证安全访问 feat(auth): 完善用户登录注册及权限相关接口与页面 - 实现手机号验证码发送及校验功能，保障注册安全 - 支持手机号注册、登录及退出接口，集成日志记录 - 增加修改密码功能，验证原密码后更新 - 提供动态导航菜单接口，基于权限展示不同菜单 - 实现管理界面路由及日志、角色、字典管理页面访问权限控制 - 添加系统日志查询接口，支持关键词和等级筛选 feat(app): 初始化Flask应用并配置蓝图与数据库 - 创建应用程序工厂，加载配置，初始化数据库和Redis客户端 - 注册认证、API及管理员蓝图，整合路由 - 根路由渲染主页模板 - 应用上下文中自动创建数据库表，保证运行环境准备完毕 feat(database): 提供数据库创建与迁移支持脚本 - 新增数据库创建脚本，支持自动检测是否已存在 - 添加数据库表初始化脚本，支持创建和删除所有表 - 实现RBAC权限初始化，包含基础权限和角色创建 - 新增字段手动修复脚本，添加用户API Key和积分字段 - 强制迁移脚本支持清理连接和修复表结构，初始化默认数据及角色分配 feat(config): 新增系统配置参数 - 配置数据库、Redis、Session和MinIO相关参数 - 添加AI接口地址及试用Key配置 - 集成阿里云短信服务配置及开发模式相关参数 feat(extensions): 初始化数据库、Redis和MinIO客户端 - 创建全局SQLAlchemy数据库实例和Redis客户端 - 配置基于boto3的MinIO兼容S3客户端 chore(logs): 添加示例系统日志文件 - 记录用户请求、验证码发送成功与失败的日志信息
765 lines
28 KiB
Python
765 lines
28 KiB
Python
# Copyright 2016 Amazon.com, Inc. or its affiliates. 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. A copy of
|
|
# the License is located at
|
|
#
|
|
# http://aws.amazon.com/apache2.0/
|
|
#
|
|
# or in the "license" file accompanying this file. This file 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 copy
|
|
import logging
|
|
import re
|
|
import threading
|
|
|
|
from s3transfer.bandwidth import BandwidthLimiter, LeakyBucket
|
|
from s3transfer.constants import (
|
|
ALLOWED_DOWNLOAD_ARGS,
|
|
FULL_OBJECT_CHECKSUM_ARGS,
|
|
KB,
|
|
MB,
|
|
)
|
|
from s3transfer.copies import CopySubmissionTask
|
|
from s3transfer.delete import DeleteSubmissionTask
|
|
from s3transfer.download import DownloadSubmissionTask
|
|
from s3transfer.exceptions import CancelledError, FatalError
|
|
from s3transfer.futures import (
|
|
IN_MEMORY_DOWNLOAD_TAG,
|
|
IN_MEMORY_UPLOAD_TAG,
|
|
BoundedExecutor,
|
|
TransferCoordinator,
|
|
TransferFuture,
|
|
TransferMeta,
|
|
)
|
|
from s3transfer.upload import UploadSubmissionTask
|
|
from s3transfer.utils import (
|
|
CallArgs,
|
|
OSUtils,
|
|
SlidingWindowSemaphore,
|
|
TaskSemaphore,
|
|
get_callbacks,
|
|
set_default_checksum_algorithm,
|
|
signal_not_transferring,
|
|
signal_transferring,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class TransferConfig:
|
|
UNSET_DEFAULT = object()
|
|
|
|
def __init__(
|
|
self,
|
|
multipart_threshold=8 * MB,
|
|
multipart_chunksize=8 * MB,
|
|
max_request_concurrency=10,
|
|
max_submission_concurrency=5,
|
|
max_request_queue_size=1000,
|
|
max_submission_queue_size=1000,
|
|
max_io_queue_size=1000,
|
|
io_chunksize=256 * KB,
|
|
num_download_attempts=5,
|
|
max_in_memory_upload_chunks=10,
|
|
max_in_memory_download_chunks=10,
|
|
max_bandwidth=None,
|
|
):
|
|
"""Configurations for the transfer manager
|
|
|
|
:param multipart_threshold: The threshold for which multipart
|
|
transfers occur.
|
|
|
|
:param max_request_concurrency: The maximum number of S3 API
|
|
transfer-related requests that can happen at a time.
|
|
|
|
:param max_submission_concurrency: The maximum number of threads
|
|
processing a call to a TransferManager method. Processing a
|
|
call usually entails determining which S3 API requests that need
|
|
to be enqueued, but does **not** entail making any of the
|
|
S3 API data transferring requests needed to perform the transfer.
|
|
The threads controlled by ``max_request_concurrency`` is
|
|
responsible for that.
|
|
|
|
:param multipart_chunksize: The size of each transfer if a request
|
|
becomes a multipart transfer.
|
|
|
|
:param max_request_queue_size: The maximum amount of S3 API requests
|
|
that can be queued at a time.
|
|
|
|
:param max_submission_queue_size: The maximum amount of
|
|
TransferManager method calls that can be queued at a time.
|
|
|
|
:param max_io_queue_size: The maximum amount of read parts that
|
|
can be queued to be written to disk per download. The default
|
|
size for each elementin this queue is 8 KB.
|
|
|
|
:param io_chunksize: The max size of each chunk in the io queue.
|
|
Currently, this is size used when reading from the downloaded
|
|
stream as well.
|
|
|
|
:param num_download_attempts: The number of download attempts that
|
|
will be tried upon errors with downloading an object in S3. Note
|
|
that these retries account for errors that occur when streaming
|
|
down the data from s3 (i.e. socket errors and read timeouts that
|
|
occur after receiving an OK response from s3).
|
|
Other retryable exceptions such as throttling errors and 5xx errors
|
|
are already retried by botocore (this default is 5). The
|
|
``num_download_attempts`` does not take into account the
|
|
number of exceptions retried by botocore.
|
|
|
|
:param max_in_memory_upload_chunks: The number of chunks that can
|
|
be stored in memory at a time for all ongoing upload requests.
|
|
This pertains to chunks of data that need to be stored in memory
|
|
during an upload if the data is sourced from a file-like object.
|
|
The total maximum memory footprint due to a in-memory upload
|
|
chunks is roughly equal to:
|
|
|
|
max_in_memory_upload_chunks * multipart_chunksize
|
|
+ max_submission_concurrency * multipart_chunksize
|
|
|
|
``max_submission_concurrency`` has an affect on this value because
|
|
for each thread pulling data off of a file-like object, they may
|
|
be waiting with a single read chunk to be submitted for upload
|
|
because the ``max_in_memory_upload_chunks`` value has been reached
|
|
by the threads making the upload request.
|
|
|
|
:param max_in_memory_download_chunks: The number of chunks that can
|
|
be buffered in memory and **not** in the io queue at a time for all
|
|
ongoing download requests. This pertains specifically to file-like
|
|
objects that cannot be seeked. The total maximum memory footprint
|
|
due to a in-memory download chunks is roughly equal to:
|
|
|
|
max_in_memory_download_chunks * multipart_chunksize
|
|
|
|
:param max_bandwidth: The maximum bandwidth that will be consumed
|
|
in uploading and downloading file content. The value is in terms of
|
|
bytes per second.
|
|
"""
|
|
self.multipart_threshold = multipart_threshold
|
|
self.multipart_chunksize = multipart_chunksize
|
|
self.max_request_concurrency = max_request_concurrency
|
|
self.max_submission_concurrency = max_submission_concurrency
|
|
self.max_request_queue_size = max_request_queue_size
|
|
self.max_submission_queue_size = max_submission_queue_size
|
|
self.max_io_queue_size = max_io_queue_size
|
|
self.io_chunksize = io_chunksize
|
|
self.num_download_attempts = num_download_attempts
|
|
self.max_in_memory_upload_chunks = max_in_memory_upload_chunks
|
|
self.max_in_memory_download_chunks = max_in_memory_download_chunks
|
|
self.max_bandwidth = max_bandwidth
|
|
self._validate_attrs_are_nonzero()
|
|
|
|
def _validate_attrs_are_nonzero(self):
|
|
for attr, attr_val in self.__dict__.items():
|
|
if (
|
|
attr_val is not None
|
|
and attr_val is not self.UNSET_DEFAULT
|
|
and attr_val <= 0
|
|
):
|
|
raise ValueError(
|
|
f'Provided parameter {attr} of value {attr_val} must '
|
|
'be greater than 0.'
|
|
)
|
|
|
|
def get_deep_attr(self, item):
|
|
return object.__getattribute__(self, item)
|
|
|
|
|
|
class TransferManager:
|
|
ALLOWED_DOWNLOAD_ARGS = ALLOWED_DOWNLOAD_ARGS
|
|
|
|
_ALLOWED_SHARED_ARGS = [
|
|
'ACL',
|
|
'CacheControl',
|
|
'ChecksumAlgorithm',
|
|
'ContentDisposition',
|
|
'ContentEncoding',
|
|
'ContentLanguage',
|
|
'ContentType',
|
|
'ExpectedBucketOwner',
|
|
'Expires',
|
|
'GrantFullControl',
|
|
'GrantRead',
|
|
'GrantReadACP',
|
|
'GrantWriteACP',
|
|
'Metadata',
|
|
'ObjectLockLegalHoldStatus',
|
|
'ObjectLockMode',
|
|
'ObjectLockRetainUntilDate',
|
|
'RequestPayer',
|
|
'ServerSideEncryption',
|
|
'StorageClass',
|
|
'SSECustomerAlgorithm',
|
|
'SSECustomerKey',
|
|
'SSECustomerKeyMD5',
|
|
'SSEKMSKeyId',
|
|
'SSEKMSEncryptionContext',
|
|
'Tagging',
|
|
'WebsiteRedirectLocation',
|
|
]
|
|
|
|
ALLOWED_UPLOAD_ARGS = (
|
|
_ALLOWED_SHARED_ARGS
|
|
+ [
|
|
'ChecksumType',
|
|
'MpuObjectSize',
|
|
]
|
|
+ FULL_OBJECT_CHECKSUM_ARGS
|
|
)
|
|
|
|
ALLOWED_COPY_ARGS = _ALLOWED_SHARED_ARGS + [
|
|
'CopySourceIfMatch',
|
|
'CopySourceIfModifiedSince',
|
|
'CopySourceIfNoneMatch',
|
|
'CopySourceIfUnmodifiedSince',
|
|
'CopySourceSSECustomerAlgorithm',
|
|
'CopySourceSSECustomerKey',
|
|
'CopySourceSSECustomerKeyMD5',
|
|
'MetadataDirective',
|
|
'TaggingDirective',
|
|
]
|
|
|
|
ALLOWED_DELETE_ARGS = [
|
|
'MFA',
|
|
'VersionId',
|
|
'RequestPayer',
|
|
'ExpectedBucketOwner',
|
|
]
|
|
|
|
VALIDATE_SUPPORTED_BUCKET_VALUES = True
|
|
|
|
_UNSUPPORTED_BUCKET_PATTERNS = {
|
|
'S3 Object Lambda': re.compile(
|
|
r'^arn:(aws).*:s3-object-lambda:[a-z\-0-9]+:[0-9]{12}:'
|
|
r'accesspoint[/:][a-zA-Z0-9\-]{1,63}'
|
|
),
|
|
}
|
|
|
|
def __init__(self, client, config=None, osutil=None, executor_cls=None):
|
|
"""A transfer manager interface for Amazon S3
|
|
|
|
:param client: Client to be used by the manager
|
|
:param config: TransferConfig to associate specific configurations
|
|
:param osutil: OSUtils object to use for os-related behavior when
|
|
using with transfer manager.
|
|
|
|
:type executor_cls: s3transfer.futures.BaseExecutor
|
|
:param executor_cls: The class of executor to use with the transfer
|
|
manager. By default, concurrent.futures.ThreadPoolExecutor is used.
|
|
"""
|
|
self._client = client
|
|
self._config = config
|
|
if config is None:
|
|
self._config = TransferConfig()
|
|
self._osutil = osutil
|
|
if osutil is None:
|
|
self._osutil = OSUtils()
|
|
self._coordinator_controller = TransferCoordinatorController()
|
|
# A counter to create unique id's for each transfer submitted.
|
|
self._id_counter = 0
|
|
|
|
# The executor responsible for making S3 API transfer requests
|
|
self._request_executor = BoundedExecutor(
|
|
max_size=self._config.max_request_queue_size,
|
|
max_num_threads=self._config.max_request_concurrency,
|
|
tag_semaphores={
|
|
IN_MEMORY_UPLOAD_TAG: TaskSemaphore(
|
|
self._config.max_in_memory_upload_chunks
|
|
),
|
|
IN_MEMORY_DOWNLOAD_TAG: SlidingWindowSemaphore(
|
|
self._config.max_in_memory_download_chunks
|
|
),
|
|
},
|
|
executor_cls=executor_cls,
|
|
)
|
|
|
|
# The executor responsible for submitting the necessary tasks to
|
|
# perform the desired transfer
|
|
self._submission_executor = BoundedExecutor(
|
|
max_size=self._config.max_submission_queue_size,
|
|
max_num_threads=self._config.max_submission_concurrency,
|
|
executor_cls=executor_cls,
|
|
)
|
|
|
|
# There is one thread available for writing to disk. It will handle
|
|
# downloads for all files.
|
|
self._io_executor = BoundedExecutor(
|
|
max_size=self._config.max_io_queue_size,
|
|
max_num_threads=1,
|
|
executor_cls=executor_cls,
|
|
)
|
|
|
|
# The component responsible for limiting bandwidth usage if it
|
|
# is configured.
|
|
self._bandwidth_limiter = None
|
|
if self._config.max_bandwidth is not None:
|
|
logger.debug(
|
|
'Setting max_bandwidth to %s', self._config.max_bandwidth
|
|
)
|
|
leaky_bucket = LeakyBucket(self._config.max_bandwidth)
|
|
self._bandwidth_limiter = BandwidthLimiter(leaky_bucket)
|
|
|
|
self._register_handlers()
|
|
|
|
@property
|
|
def client(self):
|
|
return self._client
|
|
|
|
@property
|
|
def config(self):
|
|
return self._config
|
|
|
|
def upload(self, fileobj, bucket, key, extra_args=None, subscribers=None):
|
|
"""Uploads a file to S3
|
|
|
|
:type fileobj: str or seekable file-like object
|
|
:param fileobj: The name of a file to upload or a seekable file-like
|
|
object to upload. It is recommended to use a filename because
|
|
file-like objects may result in higher memory usage.
|
|
|
|
:type bucket: str
|
|
:param bucket: The name of the bucket to upload to
|
|
|
|
:type key: str
|
|
:param key: The name of the key to upload to
|
|
|
|
:type extra_args: dict
|
|
:param extra_args: Extra arguments that may be passed to the
|
|
client operation
|
|
|
|
:type subscribers: list(s3transfer.subscribers.BaseSubscriber)
|
|
:param subscribers: The list of subscribers to be invoked in the
|
|
order provided based on the event emit during the process of
|
|
the transfer request.
|
|
|
|
:rtype: s3transfer.futures.TransferFuture
|
|
:returns: Transfer future representing the upload
|
|
"""
|
|
|
|
extra_args = extra_args.copy() if extra_args else {}
|
|
if subscribers is None:
|
|
subscribers = []
|
|
self._validate_all_known_args(extra_args, self.ALLOWED_UPLOAD_ARGS)
|
|
self._validate_if_bucket_supported(bucket)
|
|
self._add_operation_defaults(extra_args)
|
|
call_args = CallArgs(
|
|
fileobj=fileobj,
|
|
bucket=bucket,
|
|
key=key,
|
|
extra_args=extra_args,
|
|
subscribers=subscribers,
|
|
)
|
|
extra_main_kwargs = {}
|
|
if self._bandwidth_limiter:
|
|
extra_main_kwargs['bandwidth_limiter'] = self._bandwidth_limiter
|
|
return self._submit_transfer(
|
|
call_args, UploadSubmissionTask, extra_main_kwargs
|
|
)
|
|
|
|
def download(
|
|
self, bucket, key, fileobj, extra_args=None, subscribers=None
|
|
):
|
|
"""Downloads a file from S3
|
|
|
|
:type bucket: str
|
|
:param bucket: The name of the bucket to download from
|
|
|
|
:type key: str
|
|
:param key: The name of the key to download from
|
|
|
|
:type fileobj: str or seekable file-like object
|
|
:param fileobj: The name of a file to download or a seekable file-like
|
|
object to download. It is recommended to use a filename because
|
|
file-like objects may result in higher memory usage.
|
|
|
|
:type extra_args: dict
|
|
:param extra_args: Extra arguments that may be passed to the
|
|
client operation
|
|
|
|
:type subscribers: list(s3transfer.subscribers.BaseSubscriber)
|
|
:param subscribers: The list of subscribers to be invoked in the
|
|
order provided based on the event emit during the process of
|
|
the transfer request.
|
|
|
|
:rtype: s3transfer.futures.TransferFuture
|
|
:returns: Transfer future representing the download
|
|
"""
|
|
if extra_args is None:
|
|
extra_args = {}
|
|
if subscribers is None:
|
|
subscribers = []
|
|
self._validate_all_known_args(extra_args, self.ALLOWED_DOWNLOAD_ARGS)
|
|
self._validate_if_bucket_supported(bucket)
|
|
call_args = CallArgs(
|
|
bucket=bucket,
|
|
key=key,
|
|
fileobj=fileobj,
|
|
extra_args=extra_args,
|
|
subscribers=subscribers,
|
|
)
|
|
extra_main_kwargs = {'io_executor': self._io_executor}
|
|
if self._bandwidth_limiter:
|
|
extra_main_kwargs['bandwidth_limiter'] = self._bandwidth_limiter
|
|
return self._submit_transfer(
|
|
call_args, DownloadSubmissionTask, extra_main_kwargs
|
|
)
|
|
|
|
def copy(
|
|
self,
|
|
copy_source,
|
|
bucket,
|
|
key,
|
|
extra_args=None,
|
|
subscribers=None,
|
|
source_client=None,
|
|
):
|
|
"""Copies a file in S3
|
|
|
|
:type copy_source: dict
|
|
:param copy_source: The name of the source bucket, key name of the
|
|
source object, and optional version ID of the source object. The
|
|
dictionary format is:
|
|
``{'Bucket': 'bucket', 'Key': 'key', 'VersionId': 'id'}``. Note
|
|
that the ``VersionId`` key is optional and may be omitted.
|
|
|
|
:type bucket: str
|
|
:param bucket: The name of the bucket to copy to
|
|
|
|
:type key: str
|
|
:param key: The name of the key to copy to
|
|
|
|
:type extra_args: dict
|
|
:param extra_args: Extra arguments that may be passed to the
|
|
client operation
|
|
|
|
:type subscribers: a list of subscribers
|
|
:param subscribers: The list of subscribers to be invoked in the
|
|
order provided based on the event emit during the process of
|
|
the transfer request.
|
|
|
|
:type source_client: botocore or boto3 Client
|
|
:param source_client: The client to be used for operation that
|
|
may happen at the source object. For example, this client is
|
|
used for the head_object that determines the size of the copy.
|
|
If no client is provided, the transfer manager's client is used
|
|
as the client for the source object.
|
|
|
|
:rtype: s3transfer.futures.TransferFuture
|
|
:returns: Transfer future representing the copy
|
|
"""
|
|
if extra_args is None:
|
|
extra_args = {}
|
|
if subscribers is None:
|
|
subscribers = []
|
|
if source_client is None:
|
|
source_client = self._client
|
|
self._validate_all_known_args(extra_args, self.ALLOWED_COPY_ARGS)
|
|
if isinstance(copy_source, dict):
|
|
self._validate_if_bucket_supported(copy_source.get('Bucket'))
|
|
self._validate_if_bucket_supported(bucket)
|
|
call_args = CallArgs(
|
|
copy_source=copy_source,
|
|
bucket=bucket,
|
|
key=key,
|
|
extra_args=extra_args,
|
|
subscribers=subscribers,
|
|
source_client=source_client,
|
|
)
|
|
return self._submit_transfer(call_args, CopySubmissionTask)
|
|
|
|
def delete(self, bucket, key, extra_args=None, subscribers=None):
|
|
"""Delete an S3 object.
|
|
|
|
:type bucket: str
|
|
:param bucket: The name of the bucket.
|
|
|
|
:type key: str
|
|
:param key: The name of the S3 object to delete.
|
|
|
|
:type extra_args: dict
|
|
:param extra_args: Extra arguments that may be passed to the
|
|
DeleteObject call.
|
|
|
|
:type subscribers: list
|
|
:param subscribers: A list of subscribers to be invoked during the
|
|
process of the transfer request. Note that the ``on_progress``
|
|
callback is not invoked during object deletion.
|
|
|
|
:rtype: s3transfer.futures.TransferFuture
|
|
:return: Transfer future representing the deletion.
|
|
|
|
"""
|
|
if extra_args is None:
|
|
extra_args = {}
|
|
if subscribers is None:
|
|
subscribers = []
|
|
self._validate_all_known_args(extra_args, self.ALLOWED_DELETE_ARGS)
|
|
self._validate_if_bucket_supported(bucket)
|
|
call_args = CallArgs(
|
|
bucket=bucket,
|
|
key=key,
|
|
extra_args=extra_args,
|
|
subscribers=subscribers,
|
|
)
|
|
return self._submit_transfer(call_args, DeleteSubmissionTask)
|
|
|
|
def _validate_if_bucket_supported(self, bucket):
|
|
# s3 high level operations don't support some resources
|
|
# (eg. S3 Object Lambda) only direct API calls are available
|
|
# for such resources
|
|
if self.VALIDATE_SUPPORTED_BUCKET_VALUES:
|
|
for resource, pattern in self._UNSUPPORTED_BUCKET_PATTERNS.items():
|
|
match = pattern.match(bucket)
|
|
if match:
|
|
raise ValueError(
|
|
f'TransferManager methods do not support {resource} '
|
|
'resource. Use direct client calls instead.'
|
|
)
|
|
|
|
def _validate_all_known_args(self, actual, allowed):
|
|
for kwarg in actual:
|
|
if kwarg not in allowed:
|
|
raise ValueError(
|
|
"Invalid extra_args key '{}', must be one of: {}".format(
|
|
kwarg, ', '.join(allowed)
|
|
)
|
|
)
|
|
|
|
def _add_operation_defaults(self, extra_args):
|
|
if (
|
|
self.client.meta.config.request_checksum_calculation
|
|
== "when_supported"
|
|
):
|
|
set_default_checksum_algorithm(extra_args)
|
|
|
|
def _submit_transfer(
|
|
self, call_args, submission_task_cls, extra_main_kwargs=None
|
|
):
|
|
if not extra_main_kwargs:
|
|
extra_main_kwargs = {}
|
|
|
|
# Create a TransferFuture to return back to the user
|
|
transfer_future, components = self._get_future_with_components(
|
|
call_args
|
|
)
|
|
|
|
# Add any provided done callbacks to the created transfer future
|
|
# to be invoked on the transfer future being complete.
|
|
for callback in get_callbacks(transfer_future, 'done'):
|
|
components['coordinator'].add_done_callback(callback)
|
|
|
|
# Get the main kwargs needed to instantiate the submission task
|
|
main_kwargs = self._get_submission_task_main_kwargs(
|
|
transfer_future, extra_main_kwargs
|
|
)
|
|
|
|
# Submit a SubmissionTask that will submit all of the necessary
|
|
# tasks needed to complete the S3 transfer.
|
|
self._submission_executor.submit(
|
|
submission_task_cls(
|
|
transfer_coordinator=components['coordinator'],
|
|
main_kwargs=main_kwargs,
|
|
)
|
|
)
|
|
|
|
# Increment the unique id counter for future transfer requests
|
|
self._id_counter += 1
|
|
|
|
return transfer_future
|
|
|
|
def _get_future_with_components(self, call_args):
|
|
transfer_id = self._id_counter
|
|
# Creates a new transfer future along with its components
|
|
transfer_coordinator = TransferCoordinator(transfer_id=transfer_id)
|
|
# Track the transfer coordinator for transfers to manage.
|
|
self._coordinator_controller.add_transfer_coordinator(
|
|
transfer_coordinator
|
|
)
|
|
# Also make sure that the transfer coordinator is removed once
|
|
# the transfer completes so it does not stick around in memory.
|
|
transfer_coordinator.add_done_callback(
|
|
self._coordinator_controller.remove_transfer_coordinator,
|
|
transfer_coordinator,
|
|
)
|
|
components = {
|
|
'meta': TransferMeta(call_args, transfer_id=transfer_id),
|
|
'coordinator': transfer_coordinator,
|
|
}
|
|
transfer_future = TransferFuture(**components)
|
|
return transfer_future, components
|
|
|
|
def _get_submission_task_main_kwargs(
|
|
self, transfer_future, extra_main_kwargs
|
|
):
|
|
main_kwargs = {
|
|
'client': self._client,
|
|
'config': self._config,
|
|
'osutil': self._osutil,
|
|
'request_executor': self._request_executor,
|
|
'transfer_future': transfer_future,
|
|
}
|
|
main_kwargs.update(extra_main_kwargs)
|
|
return main_kwargs
|
|
|
|
def _register_handlers(self):
|
|
# Register handlers to enable/disable callbacks on uploads.
|
|
event_name = 'request-created.s3'
|
|
self._client.meta.events.register_first(
|
|
event_name,
|
|
signal_not_transferring,
|
|
unique_id='s3upload-not-transferring',
|
|
)
|
|
self._client.meta.events.register_last(
|
|
event_name, signal_transferring, unique_id='s3upload-transferring'
|
|
)
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_value, *args):
|
|
cancel = False
|
|
cancel_msg = ''
|
|
cancel_exc_type = FatalError
|
|
# If a exception was raised in the context handler, signal to cancel
|
|
# all of the inprogress futures in the shutdown.
|
|
if exc_type:
|
|
cancel = True
|
|
cancel_msg = str(exc_value)
|
|
if not cancel_msg:
|
|
cancel_msg = repr(exc_value)
|
|
# If it was a KeyboardInterrupt, the cancellation was initiated
|
|
# by the user.
|
|
if isinstance(exc_value, KeyboardInterrupt):
|
|
cancel_exc_type = CancelledError
|
|
self._shutdown(cancel, cancel_msg, cancel_exc_type)
|
|
|
|
def shutdown(self, cancel=False, cancel_msg=''):
|
|
"""Shutdown the TransferManager
|
|
|
|
It will wait till all transfers complete before it completely shuts
|
|
down.
|
|
|
|
:type cancel: boolean
|
|
:param cancel: If True, calls TransferFuture.cancel() for
|
|
all in-progress in transfers. This is useful if you want the
|
|
shutdown to happen quicker.
|
|
|
|
:type cancel_msg: str
|
|
:param cancel_msg: The message to specify if canceling all in-progress
|
|
transfers.
|
|
"""
|
|
self._shutdown(cancel, cancel, cancel_msg)
|
|
|
|
def _shutdown(self, cancel, cancel_msg, exc_type=CancelledError):
|
|
if cancel:
|
|
# Cancel all in-flight transfers if requested, before waiting
|
|
# for them to complete.
|
|
self._coordinator_controller.cancel(cancel_msg, exc_type)
|
|
try:
|
|
# Wait until there are no more in-progress transfers. This is
|
|
# wrapped in a try statement because this can be interrupted
|
|
# with a KeyboardInterrupt that needs to be caught.
|
|
self._coordinator_controller.wait()
|
|
except KeyboardInterrupt:
|
|
# If not errors were raised in the try block, the cancel should
|
|
# have no coordinators it needs to run cancel on. If there was
|
|
# an error raised in the try statement we want to cancel all of
|
|
# the inflight transfers before shutting down to speed that
|
|
# process up.
|
|
self._coordinator_controller.cancel('KeyboardInterrupt()')
|
|
raise
|
|
finally:
|
|
# Shutdown all of the executors.
|
|
self._submission_executor.shutdown()
|
|
self._request_executor.shutdown()
|
|
self._io_executor.shutdown()
|
|
|
|
|
|
class TransferCoordinatorController:
|
|
def __init__(self):
|
|
"""Abstraction to control all transfer coordinators
|
|
|
|
This abstraction allows the manager to wait for inprogress transfers
|
|
to complete and cancel all inprogress transfers.
|
|
"""
|
|
self._lock = threading.Lock()
|
|
self._tracked_transfer_coordinators = set()
|
|
|
|
@property
|
|
def tracked_transfer_coordinators(self):
|
|
"""The set of transfer coordinators being tracked"""
|
|
with self._lock:
|
|
# We return a copy because the set is mutable and if you were to
|
|
# iterate over the set, it may be changing in length due to
|
|
# additions and removals of transfer coordinators.
|
|
return copy.copy(self._tracked_transfer_coordinators)
|
|
|
|
def add_transfer_coordinator(self, transfer_coordinator):
|
|
"""Adds a transfer coordinator of a transfer to be canceled if needed
|
|
|
|
:type transfer_coordinator: s3transfer.futures.TransferCoordinator
|
|
:param transfer_coordinator: The transfer coordinator for the
|
|
particular transfer
|
|
"""
|
|
with self._lock:
|
|
self._tracked_transfer_coordinators.add(transfer_coordinator)
|
|
|
|
def remove_transfer_coordinator(self, transfer_coordinator):
|
|
"""Remove a transfer coordinator from cancellation consideration
|
|
|
|
Typically, this method is invoked by the transfer coordinator itself
|
|
to remove its self when it completes its transfer.
|
|
|
|
:type transfer_coordinator: s3transfer.futures.TransferCoordinator
|
|
:param transfer_coordinator: The transfer coordinator for the
|
|
particular transfer
|
|
"""
|
|
with self._lock:
|
|
self._tracked_transfer_coordinators.remove(transfer_coordinator)
|
|
|
|
def cancel(self, msg='', exc_type=CancelledError):
|
|
"""Cancels all inprogress transfers
|
|
|
|
This cancels the inprogress transfers by calling cancel() on all
|
|
tracked transfer coordinators.
|
|
|
|
:param msg: The message to pass on to each transfer coordinator that
|
|
gets cancelled.
|
|
|
|
:param exc_type: The type of exception to set for the cancellation
|
|
"""
|
|
for transfer_coordinator in self.tracked_transfer_coordinators:
|
|
transfer_coordinator.cancel(msg, exc_type)
|
|
|
|
def wait(self):
|
|
"""Wait until there are no more inprogress transfers
|
|
|
|
This will not stop when failures are encountered and not propagate any
|
|
of these errors from failed transfers, but it can be interrupted with
|
|
a KeyboardInterrupt.
|
|
"""
|
|
try:
|
|
transfer_coordinator = None
|
|
for transfer_coordinator in self.tracked_transfer_coordinators:
|
|
transfer_coordinator.result()
|
|
except KeyboardInterrupt:
|
|
logger.debug('Received KeyboardInterrupt in wait()')
|
|
# If Keyboard interrupt is raised while waiting for
|
|
# the result, then exit out of the wait and raise the
|
|
# exception
|
|
if transfer_coordinator:
|
|
logger.debug(
|
|
'On KeyboardInterrupt was waiting for %s',
|
|
transfer_coordinator,
|
|
)
|
|
raise
|
|
except Exception:
|
|
# A general exception could have been thrown because
|
|
# of result(). We just want to ignore this and continue
|
|
# because we at least know that the transfer coordinator
|
|
# has completed.
|
|
pass
|