more typed orm (#28577)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
This commit is contained in:
Asuka Minato
2025-11-24 22:01:46 +09:00
committed by GitHub
parent da98a38b14
commit 751ce4ec41
13 changed files with 163 additions and 117 deletions

View File

@@ -307,7 +307,7 @@ class Dataset(Base):
return f"{dify_config.VECTOR_INDEX_NAME_PREFIX}_{normalized_dataset_id}_Node"
class DatasetProcessRule(Base):
class DatasetProcessRule(Base): # bug
__tablename__ = "dataset_process_rules"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="dataset_process_rule_pkey"),
@@ -1004,7 +1004,7 @@ class DatasetKeywordTable(TypeBase):
return None
class Embedding(Base):
class Embedding(TypeBase):
__tablename__ = "embeddings"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="embedding_pkey"),
@@ -1012,12 +1012,16 @@ class Embedding(Base):
sa.Index("created_at_idx", "created_at"),
)
id = mapped_column(StringUUID, primary_key=True, default=lambda: str(uuid4()))
model_name = mapped_column(String(255), nullable=False, server_default=sa.text("'text-embedding-ada-002'"))
hash = mapped_column(String(64), nullable=False)
embedding = mapped_column(BinaryData, nullable=False)
created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
provider_name = mapped_column(String(255), nullable=False, server_default=sa.text("''"))
id: Mapped[str] = mapped_column(StringUUID, primary_key=True, default=lambda: str(uuid4()), init=False)
model_name: Mapped[str] = mapped_column(
String(255), nullable=False, server_default=sa.text("'text-embedding-ada-002'")
)
hash: Mapped[str] = mapped_column(String(64), nullable=False)
embedding: Mapped[bytes] = mapped_column(BinaryData, nullable=False)
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
provider_name: Mapped[str] = mapped_column(String(255), nullable=False, server_default=sa.text("''"))
def set_embedding(self, embedding_data: list[float]):
self.embedding = pickle.dumps(embedding_data, protocol=pickle.HIGHEST_PROTOCOL)
@@ -1214,7 +1218,7 @@ class RateLimitLog(TypeBase):
)
class DatasetMetadata(Base):
class DatasetMetadata(TypeBase):
__tablename__ = "dataset_metadatas"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="dataset_metadata_pkey"),
@@ -1222,20 +1226,26 @@ class DatasetMetadata(Base):
sa.Index("dataset_metadata_dataset_idx", "dataset_id"),
)
id = mapped_column(StringUUID, default=lambda: str(uuid4()))
tenant_id = mapped_column(StringUUID, nullable=False)
dataset_id = mapped_column(StringUUID, nullable=False)
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4()), init=False)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
type: Mapped[str] = mapped_column(String(255), nullable=False)
name: Mapped[str] = mapped_column(String(255), nullable=False)
created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=sa.func.current_timestamp())
updated_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=sa.func.current_timestamp(), onupdate=func.current_timestamp()
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=sa.func.current_timestamp(), init=False
)
created_by = mapped_column(StringUUID, nullable=False)
updated_by = mapped_column(StringUUID, nullable=True)
updated_at: Mapped[datetime] = mapped_column(
DateTime,
nullable=False,
server_default=sa.func.current_timestamp(),
onupdate=func.current_timestamp(),
init=False,
)
created_by: Mapped[str] = mapped_column(StringUUID, nullable=False)
updated_by: Mapped[str] = mapped_column(StringUUID, nullable=True, default=None)
class DatasetMetadataBinding(Base):
class DatasetMetadataBinding(TypeBase):
__tablename__ = "dataset_metadata_bindings"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="dataset_metadata_binding_pkey"),
@@ -1245,13 +1255,15 @@ class DatasetMetadataBinding(Base):
sa.Index("dataset_metadata_binding_document_idx", "document_id"),
)
id = mapped_column(StringUUID, default=lambda: str(uuid4()))
tenant_id = mapped_column(StringUUID, nullable=False)
dataset_id = mapped_column(StringUUID, nullable=False)
metadata_id = mapped_column(StringUUID, nullable=False)
document_id = mapped_column(StringUUID, nullable=False)
created_at: Mapped[datetime] = mapped_column(DateTime, nullable=False, server_default=func.current_timestamp())
created_by = mapped_column(StringUUID, nullable=False)
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuid4()), init=False)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
dataset_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
metadata_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
document_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
created_by: Mapped[str] = mapped_column(StringUUID, nullable=False)
class PipelineBuiltInTemplate(TypeBase):
@@ -1319,22 +1331,30 @@ class PipelineCustomizedTemplate(TypeBase):
return ""
class Pipeline(Base): # type: ignore[name-defined]
class Pipeline(TypeBase):
__tablename__ = "pipelines"
__table_args__ = (sa.PrimaryKeyConstraint("id", name="pipeline_pkey"),)
id = mapped_column(StringUUID, default=lambda: str(uuidv7()))
id: Mapped[str] = mapped_column(StringUUID, default=lambda: str(uuidv7()), init=False)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
name = mapped_column(sa.String(255), nullable=False)
description = mapped_column(LongText, nullable=False, default=sa.text("''"))
workflow_id = mapped_column(StringUUID, nullable=True)
is_public = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"))
is_published = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"))
created_by = mapped_column(StringUUID, nullable=True)
created_at = mapped_column(sa.DateTime, nullable=False, server_default=func.current_timestamp())
updated_by = mapped_column(StringUUID, nullable=True)
updated_at = mapped_column(
sa.DateTime, nullable=False, server_default=func.current_timestamp(), onupdate=func.current_timestamp()
name: Mapped[str] = mapped_column(sa.String(255), nullable=False)
description: Mapped[str] = mapped_column(LongText, nullable=False, default=sa.text("''"))
workflow_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None)
is_public: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"), default=False)
is_published: Mapped[bool] = mapped_column(
sa.Boolean, nullable=False, server_default=sa.text("false"), default=False
)
created_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None)
created_at: Mapped[datetime] = mapped_column(
sa.DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None)
updated_at: Mapped[datetime] = mapped_column(
sa.DateTime,
nullable=False,
server_default=func.current_timestamp(),
onupdate=func.current_timestamp(),
init=False,
)
def retrieve_dataset(self, session: Session):