Skip to content

Commit b85ba19

Browse files
committed
feat: add compression for file storage and retrieval in knowledge.py
1 parent 8a9e0ab commit b85ba19

File tree

1 file changed

+23
-2
lines changed

1 file changed

+23
-2
lines changed

apps/knowledge/models/knowledge.py

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import io
2+
import zipfile
13
from enum import Enum
24

35
import uuid_utils.compat as uuid
@@ -276,19 +278,38 @@ class Meta:
276278

277279
def save(self, bytea=None, force_insert=False, force_update=False, using=None, update_fields=None):
278280
sha256_hash = get_sha256_hash(bytea)
281+
# 创建压缩文件
282+
zip_buffer = io.BytesIO()
283+
with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
284+
# 设置压缩级别为最高(9)
285+
zipinfo = zipfile.ZipInfo(self.file_name)
286+
zipinfo.compress_type = zipfile.ZIP_DEFLATED
287+
zip_file.writestr(zipinfo, bytea, compresslevel=9)
288+
# 获取压缩后的数据
289+
compressed_data = zip_buffer.getvalue()
279290
f = QuerySet(File).filter(sha256_hash=sha256_hash).first()
280291
if f is not None:
281292
self.loid = f.loid
282293
else:
283294
result = select_one("SELECT lo_from_bytea(%s, %s::bytea) as loid", [0, bytea])
284295
self.loid = result['loid']
285-
self.file_size = len(bytea)
296+
self.file_size = len(compressed_data)
286297
self.sha256_hash = sha256_hash
298+
# 可以在元数据中记录原始大小
299+
if 'original_size' not in self.meta:
300+
self.meta['original_size'] = len(bytea)
287301
super().save()
288302

289303
def get_bytes(self):
290304
result = select_one(f'SELECT lo_get({self.loid}) as "data"', [])
291-
return result['data']
305+
compressed_data = result['data']
306+
try:
307+
# 解压数据
308+
with zipfile.ZipFile(io.BytesIO(compressed_data)) as zip_file:
309+
return zip_file.read(self.file_name)
310+
except Exception as e:
311+
# 如果数据不是zip格式,直接返回原始数据
312+
return compressed_data
292313

293314

294315
@receiver(pre_delete, sender=File)

0 commit comments

Comments
 (0)