SunnyQjm/dify
1
0 Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-10-04 13:47:55 +08:00
Code Issues Projects Releases Wiki Activity

refactor(api): establish dataset application boundaries (#41729)

Browse Source 
Co-authored-by: Jyong <jyong@dify.ai>
This commit is contained in:
Byron.wangandJyong authored and GitHub committed 2026-09-29 08:42:31 +00:00
1 parent c005932db9
commit a67cd87eb0
417 files changed
+30522 -23751
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

api/fields/dataset_fields.py
+13 -285
View File
@@ -1,26 +1,21 @@
from collections.abc import Sequence
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any
import sqlalchemy as sa
from pydantic import Field, field_validator
from sqlalchemy import func, select
from sqlalchemy.orm import Session
from pydantic import BaseModel, Field, field_validator
from fields.base import ResponseModel
from libs.helper import to_timestamp
from models.account import Account
from models.dataset import (
AppDatasetJoin,
Dataset,
DatasetMetadata,
Document,
ExternalKnowledgeApis,
ExternalKnowledgeBindings,
Pipeline,
)
from models.model import App, Tag, TagBinding
from services.data_source.entities.notion_import import NotionPageType
class NotionEstimatePagePayload(BaseModel):
page_id: str = Field(min_length=1)
page_type: NotionPageType = Field(alias="type")
class NotionEstimateWorkspacePayload(BaseModel):
workspace_id: str = Field(min_length=1)
credential_id: str = Field(min_length=1)
pages: list[NotionEstimatePagePayload] = Field(min_length=1)
class DatasetMetadataResponse(ResponseModel):
@@ -206,270 +201,3 @@ class DatasetDetailResponse(ResponseModel):
@classmethod
def _expand_null_nested(cls, value: object) -> object:
return {} if value is None else value
@dataclass(frozen=True)
class DatasetDetailPrefetch:
"""Batch-loaded values backing the session-scoped fields of a page of datasets.
Building a detail response for a single dataset issues about ten queries for
that row alone, so a list endpoint costs one query per dataset per field.
Loading the same values for every dataset on the page up front keeps the
query count bounded regardless of the page size.
"""
app_counts: dict[str, int] = field(default_factory=dict)
document_counts: dict[str, int] = field(default_factory=dict)
word_counts: dict[str, int] = field(default_factory=dict)
available_document_counts: dict[str, int] = field(default_factory=dict)
author_names: dict[str, str] = field(default_factory=dict)
tags: dict[tuple[str, str], list[Tag]] = field(default_factory=dict)
doc_forms: dict[tuple[str, str], str | None] = field(default_factory=dict)
external_knowledge_infos: dict[tuple[str, str], dict[str, Any]] = field(default_factory=dict)
doc_metadata: dict[str, list[DatasetMetadata]] = field(default_factory=dict)
pipeline_published: dict[str, bool] = field(default_factory=dict)
def build_dataset_detail_prefetch(datasets: Sequence[Dataset], *, session: Session) -> DatasetDetailPrefetch:
"""Load every session-scoped detail field for the given datasets in a fixed number of queries."""
dataset_ids = [dataset.id for dataset in datasets]
if not dataset_ids:
return DatasetDetailPrefetch()
tenant_ids = {dataset.tenant_id for dataset in datasets}
app_counts = dict(
session.execute(
select(AppDatasetJoin.dataset_id, func.count(AppDatasetJoin.id))
.where(AppDatasetJoin.dataset_id.in_(dataset_ids), App.id == AppDatasetJoin.app_id)
.group_by(AppDatasetJoin.dataset_id)
)
.tuples()
.all()
)
document_counts: dict[str, int] = {}
word_counts: dict[str, int] = {}
available_document_counts: dict[str, int] = {}
available_document_case = sa.case(
(
sa.and_(
Document.indexing_status == "completed",
Document.enabled == True,
Document.archived == False,
),
1,
),
else_=0,
)
document_rows = session.execute(
select(
Document.dataset_id,
func.count(Document.id),
func.coalesce(func.sum(Document.word_count), 0),
func.coalesce(func.sum(available_document_case), 0),
)
.where(Document.dataset_id.in_(dataset_ids))
.group_by(Document.dataset_id)
).all()
for dataset_id, document_count, word_count, available_count in document_rows:
document_counts[dataset_id] = document_count
word_counts[dataset_id] = word_count
available_document_counts[dataset_id] = available_count
# doc_form only falls back to the documents table for datasets without a
# chunk structure of their own.
doc_forms: dict[tuple[str, str], str | None] = {}
doc_form_dataset_ids = [dataset.id for dataset in datasets if not dataset.chunk_structure]
if doc_form_dataset_ids:
doc_forms = {
(dataset_id, tenant_id): doc_form
for dataset_id, tenant_id, doc_form in session.execute(
select(Document.dataset_id, Document.tenant_id, func.min(Document.doc_form))
.where(Document.dataset_id.in_(doc_form_dataset_ids), Document.tenant_id.in_(tenant_ids))
.group_by(Document.dataset_id, Document.tenant_id)
).all()
}
author_ids = {dataset.created_by for dataset in datasets if dataset.created_by}
author_names = dict(
session.execute(select(Account.id, Account.name).where(Account.id.in_(author_ids))).tuples().all()
)
tags: dict[tuple[str, str], list[Tag]] = {}
tag_rows = session.execute(
select(TagBinding.target_id, TagBinding.tenant_id, Tag)
.join(TagBinding, Tag.id == TagBinding.tag_id)
.where(
TagBinding.target_id.in_(dataset_ids),
TagBinding.tenant_id.in_(tenant_ids),
Tag.tenant_id == TagBinding.tenant_id,
Tag.type == "knowledge",
)
).all()
for target_id, tenant_id, tag in tag_rows:
tags.setdefault((target_id, tenant_id), []).append(tag)
doc_metadata: dict[str, list[DatasetMetadata]] = {}
for dataset_metadata in session.scalars(
select(DatasetMetadata).where(DatasetMetadata.dataset_id.in_(dataset_ids))
).all():
doc_metadata.setdefault(dataset_metadata.dataset_id, []).append(dataset_metadata)
pipeline_ids = {dataset.pipeline_id for dataset in datasets if dataset.pipeline_id}
pipeline_published: dict[str, bool] = {}
if pipeline_ids:
pipeline_published = dict(
session.execute(select(Pipeline.id, Pipeline.is_published).where(Pipeline.id.in_(pipeline_ids)))
.tuples()
.all()
)
return DatasetDetailPrefetch(
app_counts=app_counts,
document_counts=document_counts,
word_counts=word_counts,
available_document_counts=available_document_counts,
author_names=author_names,
tags=tags,
doc_forms=doc_forms,
external_knowledge_infos=_load_external_knowledge_infos(datasets, session=session),
doc_metadata=doc_metadata,
pipeline_published=pipeline_published,
)
def _load_external_knowledge_infos(
datasets: Sequence[Dataset], *, session: Session
) -> dict[tuple[str, str], dict[str, Any]]:
external_datasets = [dataset for dataset in datasets if dataset.provider == "external"]
if not external_datasets:
return {}
external_dataset_ids = [dataset.id for dataset in external_datasets]
external_tenant_ids = {dataset.tenant_id for dataset in external_datasets}
bindings = session.scalars(
select(ExternalKnowledgeBindings).where(
ExternalKnowledgeBindings.dataset_id.in_(external_dataset_ids),
ExternalKnowledgeBindings.tenant_id.in_(external_tenant_ids),
)
).all()
if not bindings:
return {}
apis = {
(api.id, api.tenant_id): api
for api in session.scalars(
select(ExternalKnowledgeApis).where(
ExternalKnowledgeApis.id.in_({binding.external_knowledge_api_id for binding in bindings}),
ExternalKnowledgeApis.tenant_id.in_(external_tenant_ids),
)
).all()
}
infos: dict[tuple[str, str], dict[str, Any]] = {}
for binding in bindings:
info = Dataset.build_external_knowledge_info(
binding, apis.get((binding.external_knowledge_api_id, binding.tenant_id))
)
if info is not None:
infos[(binding.dataset_id, binding.tenant_id)] = info
return infos
@dataclass(frozen=True)
class DatasetDetailResponseSource:
"""Expose session-backed dataset fields during response validation.
With a prefetch the values come from the batch load for the whole page;
without one each field is resolved against the session on access.
"""
dataset: Any
session: Session
prefetch: DatasetDetailPrefetch | None = None
@property
def _prefetch_key(self) -> tuple[str, str]:
return (self.dataset.id, self.dataset.tenant_id)
@property
def app_count(self) -> int:
if self.prefetch is not None:
return self.prefetch.app_counts.get(self.dataset.id, 0)
return self.dataset.get_app_count(session=self.session)
@property
def document_count(self) -> int:
if self.prefetch is not None:
return self.prefetch.document_counts.get(self.dataset.id, 0)
return self.dataset.get_document_count(session=self.session)
@property
def word_count(self) -> int:
if self.prefetch is not None:
return self.prefetch.word_counts.get(self.dataset.id, 0)
return self.dataset.get_word_count(session=self.session)
@property
def author_name(self) -> str | None:
if self.prefetch is not None:
return self.prefetch.author_names.get(self.dataset.created_by)
return self.dataset.get_author_name(session=self.session)
@property
def tags(self) -> Any:
if self.prefetch is not None:
return self.prefetch.tags.get(self._prefetch_key, [])
return self.dataset.get_tags(session=self.session)
@property
def doc_form(self) -> str | None:
if self.prefetch is not None:
if self.dataset.chunk_structure:
return self.dataset.chunk_structure
return self.prefetch.doc_forms.get(self._prefetch_key)
return self.dataset.get_doc_form(session=self.session)
@property
def external_knowledge_info(self) -> Any:
if self.prefetch is not None:
if self.dataset.provider != "external":
return None
return self.prefetch.external_knowledge_infos.get(self._prefetch_key)
return self.dataset.get_external_knowledge_info(session=self.session)
@property
def doc_metadata(self) -> Any:
if self.prefetch is not None:
return self.dataset.build_doc_metadata(self.prefetch.doc_metadata.get(self.dataset.id, []))
return self.dataset.get_doc_metadata(session=self.session)
@property
def is_published(self) -> bool:
if self.prefetch is not None:
if not self.dataset.pipeline_id:
return False
return self.prefetch.pipeline_published.get(self.dataset.pipeline_id, False)
return self.dataset.get_is_published(session=self.session)
@property
def total_documents(self) -> int:
if self.prefetch is not None:
return self.prefetch.document_counts.get(self.dataset.id, 0)
return self.dataset.get_total_documents(session=self.session)
@property
def total_available_documents(self) -> int:
if self.prefetch is not None:
return self.prefetch.available_document_counts.get(self.dataset.id, 0)
return self.dataset.get_total_available_documents(session=self.session)
def __getattr__(self, name: str) -> Any:
return getattr(self.dataset, name) # guard-ignore: no-new-getattr -- delegates model fields
def dataset_detail_response_source(
dataset: Any, *, session: Session, prefetch: DatasetDetailPrefetch | None = None
) -> DatasetDetailResponseSource:
return DatasetDetailResponseSource(dataset=dataset, session=session, prefetch=prefetch)