Skip to content

refactor(models): remove legacy db.session property wrappers on DocumentSegment and DatasetQuery #41794

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
asukaminato0721 merged 1 commit into from
Sep 7, 2026
Merged

refactor(models): remove legacy db.session property wrappers on DocumentSegment and DatasetQuery #41794

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 1 addition & 21 deletions api/models/dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -922,18 +922,10 @@ class DocumentSegment(TypeBase):
stopped_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None)
hit_count: Mapped[int] = mapped_column(sa.Integer, nullable=False, default=0)

@property
def dataset(self) -> Dataset | None:
return self.get_dataset(session=db.session())

def get_dataset(self, *, session: Session) -> Dataset | None:
"""Load the owning dataset with the caller-owned database session."""
return session.get(Dataset, self.dataset_id)

@property
def document(self) -> Document | None:
return self.get_document(session=db.session())

def get_document(self, *, session: Session) -> Document | None:
"""Load the owning document with the caller-owned database session."""
return session.get(Document, self.document_id)
Expand All @@ -952,10 +944,6 @@ def next_segment(self, session: Session) -> "DocumentSegment | None":
)
)

@property
def child_chunks(self):
return self.get_child_chunks(session=db.session(), include_full_doc=False)

def get_child_chunks(self, *, session: Session, include_full_doc: bool = True) -> Sequence["ChildChunk"]:
"""Load hierarchical child chunks with the caller-owned database session."""
document = session.get(Document, self.document_id)
Expand Down Expand Up @@ -1042,10 +1030,6 @@ def get_sign_content(self) -> str:

return text

@property
def attachments(self) -> list[AttachmentItem]:
return self.get_attachments(session=db.session())

def get_attachments(self, *, session: Session) -> list[AttachmentItem]:
"""Load attachment metadata with the caller-owned database session."""
# Use JOIN to fetch attachments in a single query instead of two separate queries
Expand Down Expand Up @@ -1195,10 +1179,6 @@ class DatasetQuery(TypeBase):
DateTime, nullable=False, server_default=sa.func.current_timestamp(), init=False
)

@property
def queries(self) -> list[dict[str, Any]]:
return self.get_queries(session=db.session())

def get_queries(self, *, session: Session) -> list[dict[str, Any]]:
try:
queries = json.loads(self.content)
Expand Down Expand Up @@ -1467,7 +1447,7 @@ def settings_dict(self) -> dict[str, Any] | None:
return None

@property
def dataset_bindings(self) -> list[DatasetBindingItem]:
def dataset_bindings(self):
return self.get_dataset_bindings(session=db.session())

def get_dataset_bindings(self, *, session: Session) -> list[DatasetBindingItem]:
Expand Down
4 changes: 2 additions & 2 deletions api/tests/test_containers_integration_tests/models/test_dataset_models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -322,7 +322,7 @@ def test_document_segment_dataset_property(self, db_session_with_containers: Ses
db_session_with_containers.flush()

# Act
related_dataset = segment.dataset
related_dataset = segment.get_dataset(session=db_session_with_containers)

# Assert
assert related_dataset is not None
Expand Down Expand Up @@ -369,7 +369,7 @@ def test_document_segment_document_property(self, db_session_with_containers: Se
db_session_with_containers.flush()

# Act
related_document = segment.document
related_document = segment.get_document(session=db_session_with_containers)

# Assert
assert related_document is not None
Expand Down
Loading