初始提交
This commit is contained in:
焦龙言
@@ -0,0 +1,496 @@
|
||||
import asyncio
|
||||
import ipaddress
|
||||
import os
|
||||
import socket
|
||||
import uuid
|
||||
from collections.abc import Callable
|
||||
|
||||
import psutil
|
||||
from fastapi import FastAPI, Request, applications
|
||||
from fastapi.openapi.docs import get_redoc_html, get_swagger_ui_html, get_swagger_ui_oauth2_redirect_html
|
||||
from fastapi.openapi.utils import get_openapi
|
||||
from fastapi.responses import HTMLResponse, JSONResponse
|
||||
from redis import asyncio as aioredis
|
||||
|
||||
from config.env import AppConfig
|
||||
|
||||
|
||||
class APIDocsUtil:
|
||||
"""
|
||||
API文档工具类
|
||||
"""
|
||||
|
||||
# API文档URLs
|
||||
_OPENAPI_URL = '/openapi.json'
|
||||
_PROXY_OPENAPI_URL = '/proxy-openapi.json'
|
||||
_DOCS_URL = '/docs'
|
||||
_PROXY_DOCS_URL = '/proxy-docs'
|
||||
_REDOC_URL = '/redoc'
|
||||
_PROXY_REDOC_URL = '/proxy-redoc'
|
||||
_OAUTH2_REDIRECT_URL = '/docs/oauth2-redirect'
|
||||
_PROXY_OAUTH2_REDIRECT_URL = '/proxy-docs/oauth2-redirect'
|
||||
|
||||
# 文档静态资源URLs
|
||||
DEFAULT_REDOC_JS_URL = 'https://registry.npmmirror.com/redoc/2/files/bundles/redoc.standalone.js'
|
||||
DEFAULT_REDOC_FAVICON_URL = 'https://fastapi.tiangolo.com/img/favicon.png'
|
||||
DEFAULT_SWAGGER_JS_URL = 'https://registry.npmmirror.com/swagger-ui-dist/5/files/swagger-ui-bundle.js'
|
||||
DEFAULT_SWAGGER_CSS_URL = 'https://registry.npmmirror.com/swagger-ui-dist/5/files/swagger-ui.css'
|
||||
DEFAULT_SWAGGER_FAVICON_URL = 'https://fastapi.tiangolo.com/img/favicon.png'
|
||||
|
||||
@classmethod
|
||||
def proxy_openapi_url(cls) -> str:
|
||||
"""
|
||||
代理OpenAPI文档URL
|
||||
"""
|
||||
return cls._PROXY_OPENAPI_URL if not AppConfig.app_disable_swagger and not AppConfig.app_disable_redoc else None
|
||||
|
||||
@classmethod
|
||||
def docs_url(cls) -> str:
|
||||
"""
|
||||
文档URL
|
||||
"""
|
||||
return cls._DOCS_URL
|
||||
|
||||
@classmethod
|
||||
def proxy_docs_url(cls) -> str:
|
||||
"""
|
||||
代理文档URL
|
||||
"""
|
||||
return cls._PROXY_DOCS_URL if not AppConfig.app_disable_swagger else None
|
||||
|
||||
@classmethod
|
||||
def redoc_url(cls) -> str:
|
||||
"""
|
||||
ReDoc文档URL
|
||||
"""
|
||||
return cls._REDOC_URL
|
||||
|
||||
@classmethod
|
||||
def proxy_redoc_url(cls) -> str:
|
||||
"""
|
||||
代理ReDoc文档URL
|
||||
"""
|
||||
return cls._PROXY_REDOC_URL if not AppConfig.app_disable_redoc else None
|
||||
|
||||
@classmethod
|
||||
def proxy_oauth2_redirect_url(cls) -> str:
|
||||
"""
|
||||
代理OAuth2重定向URL
|
||||
"""
|
||||
return cls._PROXY_OAUTH2_REDIRECT_URL if not AppConfig.app_disable_swagger else None
|
||||
|
||||
@classmethod
|
||||
def setup_docs_static_resources(
|
||||
cls,
|
||||
redoc_js_url: str = DEFAULT_REDOC_JS_URL,
|
||||
redoc_favicon_url: str = DEFAULT_REDOC_FAVICON_URL,
|
||||
swagger_js_url: str = DEFAULT_SWAGGER_JS_URL,
|
||||
swagger_css_url: str = DEFAULT_SWAGGER_CSS_URL,
|
||||
swagger_favicon_url: str = DEFAULT_SWAGGER_FAVICON_URL,
|
||||
) -> None:
|
||||
"""
|
||||
配置文档静态资源
|
||||
|
||||
:param redoc_js_url: 用于加载ReDoc JavaScript的URL
|
||||
:param redoc_favicon_url: ReDoc要使用的favicon的URL
|
||||
:param swagger_js_url: 用于加载Swagger UI JavaScript的URL
|
||||
:param swagger_css_url: 用于加载Swagger UI CSS的URL
|
||||
:param swagger_favicon_url: Swagger UI要使用的favicon的URL
|
||||
:return:
|
||||
"""
|
||||
|
||||
def redoc_monkey_patch(*args, **kwargs) -> HTMLResponse:
|
||||
return get_redoc_html(
|
||||
*args,
|
||||
**kwargs,
|
||||
redoc_js_url=redoc_js_url,
|
||||
redoc_favicon_url=redoc_favicon_url,
|
||||
)
|
||||
|
||||
def swagger_ui_monkey_patch(*args, **kwargs) -> HTMLResponse:
|
||||
return get_swagger_ui_html(
|
||||
*args,
|
||||
**kwargs,
|
||||
swagger_js_url=swagger_js_url,
|
||||
swagger_css_url=swagger_css_url,
|
||||
swagger_favicon_url=swagger_favicon_url,
|
||||
)
|
||||
|
||||
applications.get_redoc_html = redoc_monkey_patch
|
||||
applications.get_swagger_ui_html = swagger_ui_monkey_patch
|
||||
|
||||
@classmethod
|
||||
def custom_api_docs_router(
|
||||
cls,
|
||||
app: FastAPI,
|
||||
redoc_js_url: str = DEFAULT_REDOC_JS_URL,
|
||||
redoc_favicon_url: str = DEFAULT_REDOC_FAVICON_URL,
|
||||
swagger_js_url: str = DEFAULT_SWAGGER_JS_URL,
|
||||
swagger_css_url: str = DEFAULT_SWAGGER_CSS_URL,
|
||||
swagger_favicon_url: str = DEFAULT_SWAGGER_FAVICON_URL,
|
||||
) -> None:
|
||||
"""
|
||||
自定义API文档路由
|
||||
|
||||
:param app: FastAPI对象
|
||||
:param redoc_js_url: 用于加载ReDoc JavaScript的URL
|
||||
:param redoc_favicon_url: ReDoc要使用的favicon的URL
|
||||
:param swagger_js_url: 用于加载Swagger UI JavaScript的URL
|
||||
:param swagger_css_url: 用于加载Swagger UI CSS的URL
|
||||
:param swagger_favicon_url: Swagger UI要使用的favicon的URL
|
||||
:return:
|
||||
"""
|
||||
|
||||
async def custom_openapi(request: Request) -> JSONResponse:
|
||||
return await cls._custom_openapi(app)
|
||||
|
||||
async def custom_redoc(request: Request) -> HTMLResponse:
|
||||
return await cls._custom_redoc(app, redoc_js_url, redoc_favicon_url)
|
||||
|
||||
async def custom_swagger(request: Request) -> HTMLResponse:
|
||||
return await cls._custom_swagger(app, swagger_js_url, swagger_css_url, swagger_favicon_url)
|
||||
|
||||
async def custom_swagger_ui_redirect(request: Request) -> HTMLResponse:
|
||||
return await cls._custom_swagger_ui_redirect(app, swagger_favicon_url)
|
||||
|
||||
# 注册路由
|
||||
app.add_route(cls._OPENAPI_URL, custom_openapi, include_in_schema=False)
|
||||
cls._register_docs_routes(app, custom_swagger, custom_swagger_ui_redirect, custom_redoc)
|
||||
|
||||
@classmethod
|
||||
async def _custom_openapi(cls, app: FastAPI) -> JSONResponse:
|
||||
"""
|
||||
自定义 OpenAPI 路由处理函数
|
||||
|
||||
:param app: FastAPI对象
|
||||
:return: openapi的json响应
|
||||
"""
|
||||
openapi_schema = get_openapi(
|
||||
title=app.title,
|
||||
version=app.version,
|
||||
openapi_version=app.openapi_version,
|
||||
summary=app.summary,
|
||||
description=app.description,
|
||||
terms_of_service=app.terms_of_service,
|
||||
contact=app.contact,
|
||||
license_info=app.license_info,
|
||||
routes=app.routes,
|
||||
webhooks=app.webhooks.routes,
|
||||
tags=app.openapi_tags,
|
||||
separate_input_output_schemas=app.separate_input_output_schemas,
|
||||
external_docs=app.openapi_external_docs,
|
||||
)
|
||||
return JSONResponse(openapi_schema)
|
||||
|
||||
@classmethod
|
||||
async def _custom_redoc(
|
||||
cls,
|
||||
app: FastAPI,
|
||||
redoc_js_url: str,
|
||||
redoc_favicon_url: str,
|
||||
) -> HTMLResponse:
|
||||
"""
|
||||
自定义 ReDoc 路由处理函数
|
||||
|
||||
:param app: FastAPI对象
|
||||
:param redoc_js_url: 用于加载ReDoc JavaScript的URL
|
||||
:param redoc_favicon_url: ReDoc要使用的favicon的URL
|
||||
:return: ReDoc HTML响应
|
||||
"""
|
||||
if not AppConfig.app_disable_redoc:
|
||||
return get_redoc_html(
|
||||
openapi_url=cls._OPENAPI_URL,
|
||||
title=f'{app.title} - ReDoc',
|
||||
redoc_js_url=redoc_js_url,
|
||||
redoc_favicon_url=redoc_favicon_url,
|
||||
)
|
||||
return cls._get_disabled_html_content(
|
||||
f'{app.title} - ReDoc',
|
||||
'ReDoc',
|
||||
redoc_favicon_url,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def _custom_swagger(
|
||||
cls,
|
||||
app: FastAPI,
|
||||
swagger_js_url: str,
|
||||
swagger_css_url: str,
|
||||
swagger_favicon_url: str,
|
||||
) -> HTMLResponse:
|
||||
"""
|
||||
自定义 Swagger UI 路由处理函数
|
||||
|
||||
:param app: FastAPI对象
|
||||
:param swagger_js_url: 用于加载Swagger UI JavaScript的URL
|
||||
:param swagger_css_url: 用于加载Swagger UI CSS的URL
|
||||
:param swagger_favicon_url: Swagger UI要使用的favicon的URL
|
||||
:return: Swagger UI HTML响应
|
||||
"""
|
||||
if not AppConfig.app_disable_swagger:
|
||||
return get_swagger_ui_html(
|
||||
openapi_url=cls._OPENAPI_URL,
|
||||
title=f'{app.title} - Swagger UI',
|
||||
swagger_js_url=swagger_js_url,
|
||||
swagger_css_url=swagger_css_url,
|
||||
swagger_favicon_url=swagger_favicon_url,
|
||||
oauth2_redirect_url=app.swagger_ui_oauth2_redirect_url,
|
||||
init_oauth=app.swagger_ui_init_oauth,
|
||||
swagger_ui_parameters=app.swagger_ui_parameters,
|
||||
)
|
||||
return cls._get_disabled_html_content(
|
||||
f'{app.title} - Swagger UI',
|
||||
'Swagger UI',
|
||||
swagger_favicon_url,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def _custom_swagger_ui_redirect(
|
||||
cls,
|
||||
app: FastAPI,
|
||||
swagger_favicon_url: str,
|
||||
) -> HTMLResponse:
|
||||
"""
|
||||
自定义 Swagger UI OAuth2 重定向路由处理函数
|
||||
|
||||
:param app: FastAPI对象
|
||||
:param swagger_favicon_url: Swagger UI要使用的favicon的URL
|
||||
:return: Swagger UI OAuth2重定向HTML响应
|
||||
"""
|
||||
if not AppConfig.app_disable_swagger:
|
||||
return get_swagger_ui_oauth2_redirect_html()
|
||||
return cls._get_disabled_html_content(
|
||||
f'{app.title} - Swagger UI OAuth2 Redirect',
|
||||
'Swagger UI OAuth2 Redirect',
|
||||
swagger_favicon_url,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _get_disabled_html_content(title: str, name: str, favicon_url: str) -> HTMLResponse:
|
||||
"""
|
||||
生成禁用文档的HTML内容
|
||||
|
||||
:param title: 页面标题
|
||||
:param name: 文档名称
|
||||
:param favicon_url: 图标地址
|
||||
:return: 禁用文档HTML响应
|
||||
"""
|
||||
html = f"""
|
||||
<!DOCTYPE html>
|
||||
<html>
|
||||
<head>
|
||||
<title>{title}</title>
|
||||
<!-- needed for adaptive design -->
|
||||
<meta charset="utf-8"/>
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<link rel="shortcut icon" href="{favicon_url}">
|
||||
<!--
|
||||
{name} doesn't change outer page styles
|
||||
-->
|
||||
<style>
|
||||
body {{
|
||||
margin: 10px;
|
||||
padding: 0;
|
||||
}}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<noscript>
|
||||
{name} requires Javascript to function. Please enable it to browse the documentation.
|
||||
</noscript>
|
||||
<h1 style="color: #ff4d4f;">{name} has been disabled. Please enable it first.</h1>
|
||||
</body>
|
||||
</html>
|
||||
"""
|
||||
return HTMLResponse(html)
|
||||
|
||||
@classmethod
|
||||
def _register_docs_routes(
|
||||
cls, app: FastAPI, swagger_handler: Callable, redirect_handler: Callable, redoc_handler: Callable
|
||||
) -> None:
|
||||
"""
|
||||
注册文档路由
|
||||
|
||||
:param app: FastAPI对象
|
||||
:param swagger_handler: Swagger UI 路由处理函数
|
||||
:param redirect_handler: Swagger UI OAuth2 重定向路由处理函数
|
||||
:param redoc_handler: ReDoc 路由处理函数
|
||||
:return:
|
||||
"""
|
||||
swagger_urls: list[str] = (
|
||||
[cls._DOCS_URL] if not AppConfig.app_disable_swagger else [cls._DOCS_URL, cls._PROXY_DOCS_URL]
|
||||
)
|
||||
swagger_redirect_urls: list[str] = (
|
||||
[cls._OAUTH2_REDIRECT_URL]
|
||||
if not AppConfig.app_disable_swagger
|
||||
else [cls._OAUTH2_REDIRECT_URL, cls._PROXY_OAUTH2_REDIRECT_URL]
|
||||
)
|
||||
redoc_urls: list[str] = (
|
||||
[cls._REDOC_URL] if not AppConfig.app_disable_redoc else [cls._REDOC_URL, cls._PROXY_REDOC_URL]
|
||||
)
|
||||
|
||||
for url in swagger_urls:
|
||||
app.add_route(url, swagger_handler, include_in_schema=False)
|
||||
for url in swagger_redirect_urls:
|
||||
app.add_route(url, redirect_handler, include_in_schema=False)
|
||||
for url in redoc_urls:
|
||||
app.add_route(url, redoc_handler, include_in_schema=False)
|
||||
|
||||
|
||||
class StartupUtil:
|
||||
"""
|
||||
启动门禁工具类
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
async def acquire_startup_log_gate(
|
||||
cls, redis: aioredis.Redis, lock_key: str, worker_id: str, lock_expire_seconds: int
|
||||
) -> bool:
|
||||
"""
|
||||
获取启动日志门禁
|
||||
|
||||
:param redis: Redis连接对象
|
||||
:param lock_key: 分布式锁key
|
||||
:param worker_id: 当前worker标识
|
||||
:param lock_expire_seconds: 锁过期时间
|
||||
:return: 是否获得启动日志输出权
|
||||
"""
|
||||
acquired = await redis.set(lock_key, worker_id, nx=True, ex=lock_expire_seconds)
|
||||
if acquired:
|
||||
return True
|
||||
current_holder = await redis.get(lock_key)
|
||||
return current_holder == worker_id
|
||||
|
||||
@classmethod
|
||||
def start_lock_renewal(
|
||||
cls,
|
||||
redis: aioredis.Redis,
|
||||
lock_key: str,
|
||||
worker_id: str,
|
||||
lock_expire_seconds: int,
|
||||
interval_seconds: int,
|
||||
on_lock_lost: Callable[[], None] | None = None,
|
||||
) -> asyncio.Task:
|
||||
"""
|
||||
启动分布式锁续期任务
|
||||
|
||||
:param redis: Redis连接对象
|
||||
:param lock_key: 分布式锁key
|
||||
:param worker_id: 当前worker标识
|
||||
:param lock_expire_seconds: 锁过期时间
|
||||
:param interval_seconds: 续期间隔时间
|
||||
:param on_lock_lost: 失去锁时的回调
|
||||
:return: 异步任务对象
|
||||
"""
|
||||
|
||||
async def _loop() -> None:
|
||||
while True:
|
||||
try:
|
||||
current_holder = await redis.get(lock_key)
|
||||
if current_holder == worker_id:
|
||||
await redis.expire(lock_key, lock_expire_seconds)
|
||||
await asyncio.sleep(interval_seconds)
|
||||
continue
|
||||
if on_lock_lost:
|
||||
on_lock_lost()
|
||||
break
|
||||
except Exception:
|
||||
await asyncio.sleep(interval_seconds)
|
||||
|
||||
return asyncio.create_task(_loop())
|
||||
|
||||
|
||||
class WorkerIdUtil:
|
||||
"""
|
||||
Worker标识生成工具类
|
||||
"""
|
||||
|
||||
_worker_id: str | None = None
|
||||
|
||||
@classmethod
|
||||
def get_worker_id(cls, configured_worker_id: str | None) -> str:
|
||||
"""
|
||||
获取当前worker标识
|
||||
|
||||
:param configured_worker_id: 配置的worker标识
|
||||
:return: 当前worker标识
|
||||
"""
|
||||
if cls._worker_id:
|
||||
return cls._worker_id
|
||||
worker_id = configured_worker_id
|
||||
if not worker_id or worker_id.lower() == 'auto':
|
||||
worker_id = f'{os.getpid()}-{uuid.uuid4().hex[:6]}'
|
||||
cls._worker_id = worker_id
|
||||
return worker_id
|
||||
|
||||
|
||||
class IPUtil:
|
||||
"""
|
||||
IP工具类
|
||||
"""
|
||||
|
||||
_PREFERRED_DNS_HOSTS: tuple[str, str] = ('223.5.5.5', '8.8.8.8')
|
||||
_DNS_CONNECT_TIMEOUT = 1
|
||||
|
||||
@classmethod
|
||||
def get_local_ip(cls) -> str:
|
||||
"""
|
||||
获取本机Local IP
|
||||
"""
|
||||
try:
|
||||
for snics in psutil.net_if_addrs().values():
|
||||
for snic in snics:
|
||||
if snic.family == socket.AF_INET and snic.address.startswith('127.'):
|
||||
return snic.address
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return '127.0.0.1'
|
||||
|
||||
@classmethod
|
||||
def get_network_ips(cls) -> list[str]:
|
||||
"""
|
||||
获取本机Network IP列表
|
||||
"""
|
||||
network_ips = []
|
||||
try:
|
||||
# 获取网卡状态
|
||||
stats = psutil.net_if_stats()
|
||||
for name, snics in psutil.net_if_addrs().items():
|
||||
# 过滤掉状态为DOWN的网卡
|
||||
if name in stats and not stats[name].isup:
|
||||
continue
|
||||
|
||||
for snic in snics:
|
||||
if snic.family == socket.AF_INET:
|
||||
try:
|
||||
ip_obj = ipaddress.ip_address(snic.address)
|
||||
if ip_obj.is_loopback or ip_obj.is_link_local:
|
||||
continue
|
||||
network_ips.append(snic.address)
|
||||
except ValueError:
|
||||
continue
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 优先显示首选出站IP
|
||||
preferred_ip = None
|
||||
for dns_host in cls._PREFERRED_DNS_HOSTS:
|
||||
try:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
|
||||
s.settimeout(cls._DNS_CONNECT_TIMEOUT)
|
||||
s.connect((dns_host, 80))
|
||||
preferred_ip = s.getsockname()[0]
|
||||
break
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
if preferred_ip:
|
||||
if preferred_ip in network_ips:
|
||||
network_ips.remove(preferred_ip)
|
||||
network_ips.insert(0, preferred_ip)
|
||||
|
||||
if not network_ips:
|
||||
network_ips = ['127.0.0.1']
|
||||
|
||||
return network_ips
|
||||
Reference in New Issue
Block a user