outman/dify
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

optimize the unused dataset clean logic

Browse Source
This commit is contained in:
jyong 2024-07-17 17:30:27 +08:00
parent b44c69ddc3
commit 0885c2ad64
2 changed files with 56 additions and 28 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
api/extensions/ext_celery.py
View File

@ -51,7 +51,7 @@ def init_app(app: Flask) -> Celery:
}, },
'clean_unused_datasets_task': { 'clean_unused_datasets_task': {
'task': 'schedule.clean_unused_datasets_task.clean_unused_datasets_task', 'task': 'schedule.clean_unused_datasets_task.clean_unused_datasets_task',
'schedule': timedelta(days=day), 'schedule': timedelta(minutes=1),
} }
} }
celery_app.conf.update( celery_app.conf.update(

82
api/schedule/clean_unused_datasets_task.py
View File

@ -9,7 +9,7 @@ from configs import dify_config
from core.rag.index_processor.index_processor_factory import IndexProcessorFactory from core.rag.index_processor.index_processor_factory import IndexProcessorFactory
from extensions.ext_database import db from extensions.ext_database import db
from models.dataset import Dataset, DatasetQuery, Document from models.dataset import Dataset, DatasetQuery, Document
from sqlalchemy import func
@app.celery.task(queue='dataset') @app.celery.task(queue='dataset')
def clean_unused_datasets_task(): def clean_unused_datasets_task():
@ -20,10 +20,46 @@ def clean_unused_datasets_task():
page = 1 page = 1
while True: while True:
try: try:
datasets = db.session.query(Dataset).filter(Dataset.created_at < thirty_days_ago) \ # Subquery for counting new documents
.order_by(Dataset.created_at.desc()).paginate(page=page, per_page=50) document_subquery_new = db.session.query(
Document.dataset_id,
func.count(Document.id).label('document_count')
).filter(
Document.indexing_status == 'completed',
Document.enabled == True,
Document.archived == False,
Document.updated_at > thirty_days_ago
).group_by(Document.dataset_id).subquery()
# Subquery for counting old documents
document_subquery_old = db.session.query(
Document.dataset_id,
func.count(Document.id).label('document_count')
).filter(
Document.indexing_status == 'completed',
Document.enabled == True,
Document.archived == False,
Document.updated_at < thirty_days_ago
).group_by(Document.dataset_id).subquery()
# Main query with join and filter
datasets = (db.session.query(Dataset)
.outerjoin(
document_subquery_new, Dataset.id == document_subquery_new.c.dataset_id
).outerjoin(
document_subquery_old, Dataset.id == document_subquery_old.c.dataset_id
).filter(
Dataset.created_at < thirty_days_ago,
func.coalesce(document_subquery_new.c.document_count, 0) == 0,
func.coalesce(document_subquery_old.c.document_count, 0) > 0
).order_by(
Dataset.created_at.desc()
).paginate(page=page, per_page=50))
except NotFound: except NotFound:
break break
if datasets.items is None or len(datasets.items) == 0:
break
page += 1 page += 1
for dataset in datasets: for dataset in datasets:
dataset_query = db.session.query(DatasetQuery).filter( dataset_query = db.session.query(DatasetQuery).filter(
@ -31,31 +67,23 @@ def clean_unused_datasets_task():
DatasetQuery.dataset_id == dataset.id DatasetQuery.dataset_id == dataset.id
).all() ).all()
if not dataset_query or len(dataset_query) == 0: if not dataset_query or len(dataset_query) == 0:
documents = db.session.query(Document).filter( try:
Document.dataset_id == dataset.id, # remove index
Document.indexing_status == 'completed', index_processor = IndexProcessorFactory(dataset.doc_form).init_index_processor()
Document.enabled == True, index_processor.clean(dataset, None)
Document.archived == False,
Document.updated_at > thirty_days_ago
).all()
if not documents or len(documents) == 0:
try:
# remove index
index_processor = IndexProcessorFactory(dataset.doc_form).init_index_processor()
index_processor.clean(dataset, None)
# update document # update document
update_params = { update_params = {
Document.enabled: False Document.enabled: False
} }
Document.query.filter_by(dataset_id=dataset.id).update(update_params) Document.query.filter_by(dataset_id=dataset.id).update(update_params)
db.session.commit() db.session.commit()
click.echo(click.style('Cleaned unused dataset {} from db success!'.format(dataset.id), click.echo(click.style('Cleaned unused dataset {} from db success!'.format(dataset.id),
fg='green')) fg='green'))
except Exception as e: except Exception as e:
click.echo( click.echo(
click.style('clean dataset index error: {} {}'.format(e.__class__.__name__, str(e)), click.style('clean dataset index error: {} {}'.format(e.__class__.__name__, str(e)),
fg='red')) fg='red'))
end_at = time.perf_counter() end_at = time.perf_counter()
click.echo(click.style('Cleaned unused dataset from db success latency: {}'.format(end_at - start_at), fg='green')) click.echo(click.style('Cleaned unused dataset from db success latency: {}'.format(end_at - start_at), fg='green'))