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): 添加示例系统日志文件 - 记录用户请求、验证码发送成功与失败的日志信息
391 lines
16 KiB
Python
391 lines
16 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
|
|
|
|
from s3transfer.utils import get_callbacks
|
|
|
|
try:
|
|
from botocore.context import start_as_current_context
|
|
except ImportError:
|
|
from contextlib import nullcontext as start_as_current_context
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Task:
|
|
"""A task associated to a TransferFuture request
|
|
|
|
This is a base class for other classes to subclass from. All subclassed
|
|
classes must implement the main() method.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
transfer_coordinator,
|
|
main_kwargs=None,
|
|
pending_main_kwargs=None,
|
|
done_callbacks=None,
|
|
is_final=False,
|
|
):
|
|
"""
|
|
:type transfer_coordinator: s3transfer.futures.TransferCoordinator
|
|
:param transfer_coordinator: The context associated to the
|
|
TransferFuture for which this Task is associated with.
|
|
|
|
:type main_kwargs: dict
|
|
:param main_kwargs: The keyword args that can be immediately supplied
|
|
to the _main() method of the task
|
|
|
|
:type pending_main_kwargs: dict
|
|
:param pending_main_kwargs: The keyword args that are depended upon
|
|
by the result from a dependent future(s). The result returned by
|
|
the future(s) will be used as the value for the keyword argument
|
|
when _main() is called. The values for each key can be:
|
|
* a single future - Once completed, its value will be the
|
|
result of that single future
|
|
* a list of futures - Once all of the futures complete, the
|
|
value used will be a list of each completed future result
|
|
value in order of when they were originally supplied.
|
|
|
|
:type done_callbacks: list of callbacks
|
|
:param done_callbacks: A list of callbacks to call once the task is
|
|
done completing. Each callback will be called with no arguments
|
|
and will be called no matter if the task succeeds or an exception
|
|
is raised.
|
|
|
|
:type is_final: boolean
|
|
:param is_final: True, to indicate that this task is the final task
|
|
for the TransferFuture request. By setting this value to True, it
|
|
will set the result of the entire TransferFuture to the result
|
|
returned by this task's main() method.
|
|
"""
|
|
self._transfer_coordinator = transfer_coordinator
|
|
|
|
self._main_kwargs = main_kwargs
|
|
if self._main_kwargs is None:
|
|
self._main_kwargs = {}
|
|
|
|
self._pending_main_kwargs = pending_main_kwargs
|
|
if pending_main_kwargs is None:
|
|
self._pending_main_kwargs = {}
|
|
|
|
self._done_callbacks = done_callbacks
|
|
if self._done_callbacks is None:
|
|
self._done_callbacks = []
|
|
|
|
self._is_final = is_final
|
|
|
|
def __repr__(self):
|
|
# These are the general main_kwarg parameters that we want to
|
|
# display in the repr.
|
|
params_to_display = [
|
|
'bucket',
|
|
'key',
|
|
'part_number',
|
|
'final_filename',
|
|
'transfer_future',
|
|
'offset',
|
|
'extra_args',
|
|
]
|
|
main_kwargs_to_display = self._get_kwargs_with_params_to_include(
|
|
self._main_kwargs, params_to_display
|
|
)
|
|
return f'{self.__class__.__name__}(transfer_id={self._transfer_coordinator.transfer_id}, {main_kwargs_to_display})'
|
|
|
|
@property
|
|
def transfer_id(self):
|
|
"""The id for the transfer request that the task belongs to"""
|
|
return self._transfer_coordinator.transfer_id
|
|
|
|
def _get_kwargs_with_params_to_include(self, kwargs, include):
|
|
filtered_kwargs = {}
|
|
for param in include:
|
|
if param in kwargs:
|
|
filtered_kwargs[param] = kwargs[param]
|
|
return filtered_kwargs
|
|
|
|
def _get_kwargs_with_params_to_exclude(self, kwargs, exclude):
|
|
filtered_kwargs = {}
|
|
for param, value in kwargs.items():
|
|
if param in exclude:
|
|
continue
|
|
filtered_kwargs[param] = value
|
|
return filtered_kwargs
|
|
|
|
def __call__(self, ctx=None):
|
|
"""The callable to use when submitting a Task to an executor"""
|
|
with start_as_current_context(ctx):
|
|
try:
|
|
# Wait for all of futures this task depends on.
|
|
self._wait_on_dependent_futures()
|
|
# Gather up all of the main keyword arguments for main().
|
|
# This includes the immediately provided main_kwargs and
|
|
# the values for pending_main_kwargs that source from the return
|
|
# values from the task's dependent futures.
|
|
kwargs = self._get_all_main_kwargs()
|
|
# If the task is not done (really only if some other related
|
|
# task to the TransferFuture had failed) then execute the task's
|
|
# main() method.
|
|
if not self._transfer_coordinator.done():
|
|
return self._execute_main(kwargs)
|
|
except Exception as e:
|
|
self._log_and_set_exception(e)
|
|
finally:
|
|
# Run any done callbacks associated to the task no matter what.
|
|
for done_callback in self._done_callbacks:
|
|
done_callback()
|
|
|
|
if self._is_final:
|
|
# If this is the final task announce that it is done if results
|
|
# are waiting on its completion.
|
|
self._transfer_coordinator.announce_done()
|
|
|
|
def _execute_main(self, kwargs):
|
|
# Do not display keyword args that should not be printed, especially
|
|
# if they are going to make the logs hard to follow.
|
|
params_to_exclude = ['data']
|
|
kwargs_to_display = self._get_kwargs_with_params_to_exclude(
|
|
kwargs, params_to_exclude
|
|
)
|
|
# Log what is about to be executed.
|
|
logger.debug(f"Executing task {self} with kwargs {kwargs_to_display}")
|
|
|
|
return_value = self._main(**kwargs)
|
|
# If the task is the final task, then set the TransferFuture's
|
|
# value to the return value from main().
|
|
if self._is_final:
|
|
self._transfer_coordinator.set_result(return_value)
|
|
return return_value
|
|
|
|
def _log_and_set_exception(self, exception):
|
|
# If an exception is ever thrown than set the exception for the
|
|
# entire TransferFuture.
|
|
logger.debug("Exception raised.", exc_info=True)
|
|
self._transfer_coordinator.set_exception(exception)
|
|
|
|
def _main(self, **kwargs):
|
|
"""The method that will be ran in the executor
|
|
|
|
This method must be implemented by subclasses from Task. main() can
|
|
be implemented with any arguments decided upon by the subclass.
|
|
"""
|
|
raise NotImplementedError('_main() must be implemented')
|
|
|
|
def _wait_on_dependent_futures(self):
|
|
# Gather all of the futures into that main() depends on.
|
|
futures_to_wait_on = []
|
|
for _, future in self._pending_main_kwargs.items():
|
|
# If the pending main keyword arg is a list then extend the list.
|
|
if isinstance(future, list):
|
|
futures_to_wait_on.extend(future)
|
|
# If the pending main keyword arg is a future append it to the list.
|
|
else:
|
|
futures_to_wait_on.append(future)
|
|
# Now wait for all of the futures to complete.
|
|
self._wait_until_all_complete(futures_to_wait_on)
|
|
|
|
def _wait_until_all_complete(self, futures):
|
|
# This is a basic implementation of the concurrent.futures.wait()
|
|
#
|
|
# concurrent.futures.wait() is not used instead because of this
|
|
# reported issue: https://bugs.python.org/issue20319.
|
|
# The issue would occasionally cause multipart uploads to hang
|
|
# when wait() was called. With this approach, it avoids the
|
|
# concurrency bug by removing any association with concurrent.futures
|
|
# implementation of waiters.
|
|
logger.debug(
|
|
'%s about to wait for the following futures %s', self, futures
|
|
)
|
|
for future in futures:
|
|
try:
|
|
logger.debug('%s about to wait for %s', self, future)
|
|
future.result()
|
|
except Exception:
|
|
# result() can also produce exceptions. We want to ignore
|
|
# these to be deferred to error handling down the road.
|
|
pass
|
|
logger.debug('%s done waiting for dependent futures', self)
|
|
|
|
def _get_all_main_kwargs(self):
|
|
# Copy over all of the kwargs that we know is available.
|
|
kwargs = copy.copy(self._main_kwargs)
|
|
|
|
# Iterate through the kwargs whose values are pending on the result
|
|
# of a future.
|
|
for key, pending_value in self._pending_main_kwargs.items():
|
|
# If the value is a list of futures, iterate though the list
|
|
# appending on the result from each future.
|
|
if isinstance(pending_value, list):
|
|
result = []
|
|
for future in pending_value:
|
|
result.append(future.result())
|
|
# Otherwise if the pending_value is a future, just wait for it.
|
|
else:
|
|
result = pending_value.result()
|
|
# Add the retrieved value to the kwargs to be sent to the
|
|
# main() call.
|
|
kwargs[key] = result
|
|
return kwargs
|
|
|
|
|
|
class SubmissionTask(Task):
|
|
"""A base class for any submission task
|
|
|
|
Submission tasks are the top-level task used to submit a series of tasks
|
|
to execute a particular transfer.
|
|
"""
|
|
|
|
def _main(self, transfer_future, **kwargs):
|
|
"""
|
|
:type transfer_future: s3transfer.futures.TransferFuture
|
|
:param transfer_future: The transfer future associated with the
|
|
transfer request that tasks are being submitted for
|
|
|
|
:param kwargs: Any additional kwargs that you may want to pass
|
|
to the _submit() method
|
|
"""
|
|
try:
|
|
self._transfer_coordinator.set_status_to_queued()
|
|
|
|
# Before submitting any tasks, run all of the on_queued callbacks
|
|
on_queued_callbacks = get_callbacks(transfer_future, 'queued')
|
|
for on_queued_callback in on_queued_callbacks:
|
|
on_queued_callback()
|
|
|
|
# Once callbacks have been ran set the status to running.
|
|
self._transfer_coordinator.set_status_to_running()
|
|
|
|
# Call the submit method to start submitting tasks to execute the
|
|
# transfer.
|
|
self._submit(transfer_future=transfer_future, **kwargs)
|
|
except BaseException as e:
|
|
# If there was an exception raised during the submission of task
|
|
# there is a chance that the final task that signals if a transfer
|
|
# is done and too run the cleanup may never have been submitted in
|
|
# the first place so we need to account accordingly.
|
|
#
|
|
# Note that BaseException is caught, instead of Exception, because
|
|
# for some implementations of executors, specifically the serial
|
|
# implementation, the SubmissionTask is directly exposed to
|
|
# KeyboardInterupts and so needs to cleanup and signal done
|
|
# for those as well.
|
|
|
|
# Set the exception, that caused the process to fail.
|
|
self._log_and_set_exception(e)
|
|
|
|
# Wait for all possibly associated futures that may have spawned
|
|
# from this submission task have finished before we announce the
|
|
# transfer done.
|
|
self._wait_for_all_submitted_futures_to_complete()
|
|
|
|
# Announce the transfer as done, which will run any cleanups
|
|
# and done callbacks as well.
|
|
self._transfer_coordinator.announce_done()
|
|
|
|
def _submit(self, transfer_future, **kwargs):
|
|
"""The submission method to be implemented
|
|
|
|
:type transfer_future: s3transfer.futures.TransferFuture
|
|
:param transfer_future: The transfer future associated with the
|
|
transfer request that tasks are being submitted for
|
|
|
|
:param kwargs: Any additional keyword arguments you want to be passed
|
|
in
|
|
"""
|
|
raise NotImplementedError('_submit() must be implemented')
|
|
|
|
def _wait_for_all_submitted_futures_to_complete(self):
|
|
# We want to wait for all futures that were submitted to
|
|
# complete as we do not want the cleanup callbacks or done callbacks
|
|
# to be called to early. The main problem is any task that was
|
|
# submitted may have submitted even more during its process and so
|
|
# we need to account accordingly.
|
|
|
|
# First get all of the futures that were submitted up to this point.
|
|
submitted_futures = self._transfer_coordinator.associated_futures
|
|
while submitted_futures:
|
|
# Wait for those futures to complete.
|
|
self._wait_until_all_complete(submitted_futures)
|
|
# However, more futures may have been submitted as we waited so
|
|
# we need to check again for any more associated futures.
|
|
possibly_more_submitted_futures = (
|
|
self._transfer_coordinator.associated_futures
|
|
)
|
|
# If the current list of submitted futures is equal to the
|
|
# the list of associated futures for when after the wait completes,
|
|
# we can ensure no more futures were submitted in waiting on
|
|
# the current list of futures to complete ultimately meaning all
|
|
# futures that may have spawned from the original submission task
|
|
# have completed.
|
|
if submitted_futures == possibly_more_submitted_futures:
|
|
break
|
|
submitted_futures = possibly_more_submitted_futures
|
|
|
|
|
|
class CreateMultipartUploadTask(Task):
|
|
"""Task to initiate a multipart upload"""
|
|
|
|
def _main(self, client, bucket, key, extra_args):
|
|
"""
|
|
:param client: The client to use when calling CreateMultipartUpload
|
|
:param bucket: The name of the bucket to upload to
|
|
:param key: The name of the key to upload to
|
|
:param extra_args: A dictionary of any extra arguments that may be
|
|
used in the initialization.
|
|
|
|
:returns: The upload id of the multipart upload
|
|
"""
|
|
# Create the multipart upload.
|
|
response = client.create_multipart_upload(
|
|
Bucket=bucket, Key=key, **extra_args
|
|
)
|
|
upload_id = response['UploadId']
|
|
|
|
# Add a cleanup if the multipart upload fails at any point.
|
|
self._transfer_coordinator.add_failure_cleanup(
|
|
client.abort_multipart_upload,
|
|
Bucket=bucket,
|
|
Key=key,
|
|
UploadId=upload_id,
|
|
)
|
|
return upload_id
|
|
|
|
|
|
class CompleteMultipartUploadTask(Task):
|
|
"""Task to complete a multipart upload"""
|
|
|
|
def _main(self, client, bucket, key, upload_id, parts, extra_args):
|
|
"""
|
|
:param client: The client to use when calling CompleteMultipartUpload
|
|
:param bucket: The name of the bucket to upload to
|
|
:param key: The name of the key to upload to
|
|
:param upload_id: The id of the upload
|
|
:param parts: A list of parts to use to complete the multipart upload::
|
|
|
|
[{'Etag': etag_value, 'PartNumber': part_number}, ...]
|
|
|
|
Each element in the list consists of a return value from
|
|
``UploadPartTask.main()``.
|
|
:param extra_args: A dictionary of any extra arguments that may be
|
|
used in completing the multipart transfer.
|
|
"""
|
|
client.complete_multipart_upload(
|
|
Bucket=bucket,
|
|
Key=key,
|
|
UploadId=upload_id,
|
|
MultipartUpload={'Parts': parts},
|
|
**extra_args,
|
|
)
|