2021-11-26 03:12:53 +00:00
|
|
|
|
from contextlib import contextmanager
|
|
|
|
|
|
2024-04-18 20:37:03 +00:00
|
|
|
|
from django.db import connections, transaction
|
2023-07-31 10:31:11 +00:00
|
|
|
|
from django.utils.encoding import force_str
|
2020-08-16 15:08:58 +00:00
|
|
|
|
|
2023-07-31 09:39:30 +00:00
|
|
|
|
from common.utils import get_logger, signer, crypto
|
2021-11-26 03:12:53 +00:00
|
|
|
|
|
2020-08-16 15:08:58 +00:00
|
|
|
|
logger = get_logger(__file__)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def get_object_if_need(model, pk):
|
|
|
|
|
if not isinstance(pk, model):
|
|
|
|
|
try:
|
|
|
|
|
return model.objects.get(id=pk)
|
|
|
|
|
except model.DoesNotExist as e:
|
|
|
|
|
logger.error(f'DoesNotExist: <{model.__name__}:{pk}> not exist')
|
|
|
|
|
raise e
|
|
|
|
|
return pk
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def get_objects_if_need(model, pks):
|
|
|
|
|
if not pks:
|
|
|
|
|
return pks
|
|
|
|
|
if not isinstance(pks[0], model):
|
|
|
|
|
objs = list(model.objects.filter(id__in=pks))
|
|
|
|
|
if len(objs) != len(pks):
|
|
|
|
|
pks = set(pks)
|
|
|
|
|
exists_pks = {o.id for o in objs}
|
|
|
|
|
not_found_pks = ','.join(pks - exists_pks)
|
|
|
|
|
logger.error(f'DoesNotExist: <{model.__name__}: {not_found_pks}>')
|
|
|
|
|
return objs
|
|
|
|
|
return pks
|
2020-10-26 11:52:50 +00:00
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def get_objects(model, pks):
|
|
|
|
|
if not pks:
|
|
|
|
|
return pks
|
|
|
|
|
|
|
|
|
|
objs = list(model.objects.filter(id__in=pks))
|
|
|
|
|
if len(objs) != len(pks):
|
|
|
|
|
pks = set(pks)
|
|
|
|
|
exists_pks = {o.id for o in objs}
|
2020-12-27 15:00:45 +00:00
|
|
|
|
not_found_pks = pks - exists_pks
|
2020-10-26 11:52:50 +00:00
|
|
|
|
logger.error(f'DoesNotExist: <{model.__name__}: {not_found_pks}>')
|
|
|
|
|
return objs
|
2021-11-18 10:47:01 +00:00
|
|
|
|
|
|
|
|
|
|
2022-01-17 03:07:27 +00:00
|
|
|
|
# 复制 django.db.close_old_connections, 因为它没有导出,ide 提示有问题
|
2021-11-18 10:47:01 +00:00
|
|
|
|
def close_old_connections():
|
|
|
|
|
for conn in connections.all():
|
|
|
|
|
conn.close_if_unusable_or_obsolete()
|
2021-11-26 03:12:53 +00:00
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@contextmanager
|
|
|
|
|
def safe_db_connection():
|
|
|
|
|
close_old_connections()
|
|
|
|
|
yield
|
|
|
|
|
close_old_connections()
|
2023-07-31 09:39:30 +00:00
|
|
|
|
|
|
|
|
|
|
2024-04-18 20:37:03 +00:00
|
|
|
|
@contextmanager
|
|
|
|
|
def open_db_connection(alias='default'):
|
|
|
|
|
connection = transaction.get_connection(alias)
|
|
|
|
|
try:
|
|
|
|
|
connection.connect()
|
|
|
|
|
with transaction.atomic():
|
|
|
|
|
yield connection
|
|
|
|
|
finally:
|
|
|
|
|
connection.close()
|
|
|
|
|
|
|
|
|
|
|
2023-07-31 09:39:30 +00:00
|
|
|
|
class Encryptor:
|
|
|
|
|
def __init__(self, value):
|
2023-07-31 10:31:11 +00:00
|
|
|
|
self.value = force_str(value)
|
2023-07-31 09:39:30 +00:00
|
|
|
|
|
|
|
|
|
def decrypt(self):
|
|
|
|
|
plain_value = crypto.decrypt(self.value)
|
|
|
|
|
|
|
|
|
|
# 如果没有解开,使用原来的signer解密
|
|
|
|
|
if not plain_value:
|
|
|
|
|
plain_value = signer.unsign(self.value) or ""
|
|
|
|
|
return plain_value
|
|
|
|
|
|
|
|
|
|
def encrypt(self):
|
|
|
|
|
return crypto.encrypt(self.value)
|