refactor: migrate SegmentAttachmentBinding to TypeBase (#34810)

This commit is contained in:
sxxtony
2026-04-12 22:22:43 -07:00
committed by GitHub
parent 3b9aad2ba7
commit 7dd507af04
3 changed files with 32 additions and 4 deletions

View File

@ -244,7 +244,7 @@ class DatasetDocumentStore:
return document_segment
def add_multimodel_documents_binding(self, segment_id: str, multimodel_documents: list[AttachmentDocument] | None):
if multimodel_documents:
if multimodel_documents and self._document_id is not None:
for multimodel_document in multimodel_documents:
binding = SegmentAttachmentBinding(
tenant_id=self._dataset.tenant_id,

View File

@ -1688,7 +1688,7 @@ class PipelineRecommendedPlugin(TypeBase):
)
class SegmentAttachmentBinding(Base):
class SegmentAttachmentBinding(TypeBase):
__tablename__ = "segment_attachment_bindings"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="segment_attachment_binding_pkey"),
@ -1701,13 +1701,17 @@ class SegmentAttachmentBinding(Base):
),
sa.Index("segment_attachment_binding_attachment_idx", "attachment_id"),
)
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuidv7()))
id: Mapped[str] = mapped_column(
StringUUID, insert_default=lambda: str(uuidv7()), default_factory=lambda: str(uuidv7()), init=False
)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
document_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
segment_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
attachment_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
created_at: Mapped[datetime] = mapped_column(sa.DateTime, nullable=False, server_default=func.current_timestamp())
created_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
class DocumentSegmentSummary(Base):

View File

@ -721,6 +721,30 @@ class TestDatasetDocumentStoreMultimodelBinding:
mock_db.session.add.assert_not_called()
def test_add_multimodel_documents_binding_with_none_document_id(self):
"""Test that no bindings are added when document_id is None."""
mock_dataset = MagicMock(spec=Dataset)
mock_dataset.id = "test-dataset-id"
mock_dataset.tenant_id = "tenant-1"
mock_attachment = MagicMock(spec=AttachmentDocument)
mock_attachment.metadata = {"doc_id": "attachment-1"}
with patch("core.rag.docstore.dataset_docstore.db") as mock_db:
mock_session = MagicMock()
mock_db.session = mock_session
store = DatasetDocumentStore(
dataset=mock_dataset,
user_id="test-user-id",
document_id=None,
)
store.add_multimodel_documents_binding("seg-1", [mock_attachment])
mock_db.session.add.assert_not_called()
class TestDatasetDocumentStoreAddDocumentsUpdateChild:
"""Tests for add_documents when updating existing documents with children."""