e_dataset.go 1.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071
  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. FileCount: a.FileCount,
  20. Status: a.Status,
  21. CreatorID: a.CreatorId,
  22. RagDataId: a.RagDataId,
  23. Sequence: &a.Sequence,
  24. }
  25. return item
  26. }
  27. // ToSchemaDataset 转换为Dataset对象
  28. func (a Dataset) ToSchemaDataset() *schema.Dataset {
  29. item := &schema.Dataset{
  30. RecordID: a.RecordID,
  31. Name: a.Name,
  32. OrgId: a.OrgId,
  33. FileCount: a.FileCount,
  34. Status: a.Status,
  35. CreatorId: a.CreatorID,
  36. RagDataId: a.RagDataId,
  37. CreatedAt: a.CreatedAt,
  38. Sequence: *a.Sequence,
  39. }
  40. return item
  41. }
  42. // Dataset Dataset实体
  43. type Dataset struct {
  44. gorm.Model
  45. RecordID string `gorm:"column:record_id;size:32;index;"`
  46. Name string `gorm:"column:name;size:30;index;"`
  47. OrgId string `gorm:"column:org_id;size:32;index;"`
  48. FileCount int `gorm:"column:file_count;"`
  49. Status bool `gorm:"column:status;"`
  50. CreatorID string `gorm:"column:creator_id;size:32;"`
  51. RagDataId string `gorm:"column:rag_data_id;size:32;"`
  52. Sequence *int `gorm:"column:sequence;auto_increment;"`
  53. }
  54. // Datasets Dataset列表
  55. type Datasets []*Dataset
  56. // ToSchemaDatasets 转换为Dataset对象列表
  57. func (a Datasets) ToSchemaDatasets() []*schema.Dataset {
  58. list := make([]*schema.Dataset, len(a))
  59. for i, item := range a {
  60. list[i] = item.ToSchemaDataset()
  61. }
  62. return list
  63. }