e_dataset.go 1.9 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374
  1. package entity
  2. import (
  3. "context"
  4. "gorm.io/gorm"
  5. "yx-dataset-server/app/schema"
  6. )
  7. // GetDatasetDB 获取Dataset存储
  8. func GetDatasetDB(ctx context.Context, defDB *gorm.DB) *gorm.DB {
  9. return getDBWithModel(ctx, defDB, Dataset{})
  10. }
  11. // SchemaDataset Dataset对象
  12. type SchemaDataset schema.Dataset
  13. // ToDataset 转换为Dataset实体
  14. func (a SchemaDataset) ToDataset() *Dataset {
  15. item := &Dataset{
  16. RecordID: a.RecordID,
  17. Name: a.Name,
  18. OrgId: a.OrgId,
  19. Type: a.Type,
  20. FileCount: a.FileCount,
  21. Status: a.Status,
  22. CreatorID: a.CreatorId,
  23. RagDataId: a.RagDataId,
  24. Sequence: &a.Sequence,
  25. }
  26. return item
  27. }
  28. // ToSchemaDataset 转换为Dataset对象
  29. func (a Dataset) ToSchemaDataset() *schema.Dataset {
  30. item := &schema.Dataset{
  31. RecordID: a.RecordID,
  32. Name: a.Name,
  33. OrgId: a.OrgId,
  34. Type: a.Type,
  35. FileCount: a.FileCount,
  36. Status: a.Status,
  37. CreatorId: a.CreatorID,
  38. RagDataId: a.RagDataId,
  39. CreatedAt: a.CreatedAt,
  40. Sequence: *a.Sequence,
  41. }
  42. return item
  43. }
  44. // Dataset Dataset实体
  45. type Dataset struct {
  46. gorm.Model
  47. RecordID string `gorm:"column:record_id;size:32;index;"`
  48. Name string `gorm:"column:name;size:30;index;"`
  49. OrgId string `gorm:"column:org_id;size:32;index;"` // 创建者所属组织(共享知识库可为空)
  50. Type int `gorm:"column:type;index;"` // 1 公共/共享 2 企业 3 个人
  51. FileCount int `gorm:"column:file_count;"`
  52. Status bool `gorm:"column:status;"`
  53. CreatorID string `gorm:"column:creator_id;size:32;"`
  54. RagDataId string `gorm:"column:rag_data_id;size:32;"`
  55. Sequence *int `gorm:"column:sequence;auto_increment;"`
  56. }
  57. // Datasets Dataset列表
  58. type Datasets []*Dataset
  59. // ToSchemaDatasets 转换为Dataset对象列表
  60. func (a Datasets) ToSchemaDatasets() []*schema.Dataset {
  61. list := make([]*schema.Dataset, len(a))
  62. for i, item := range a {
  63. list[i] = item.ToSchemaDataset()
  64. }
  65. return list
  66. }