From a35caf03d4dd4fa1ddd4e9d37196d5f5b5bfd1ec Mon Sep 17 00:00:00 2001 From: James Kwon <96548424+hongil0316@users.noreply.github.com> Date: Wed, 8 Jan 2025 12:36:18 -0500 Subject: [PATCH] store comfy node extraction status --- drip/api.gen.go | 76 ++++++------- ent/migrate/schema.go | 7 +- ent/mutation.go | 102 +++++++++++++----- ent/nodeversion.go | 13 ++- ent/nodeversion/nodeversion.go | 10 ++ ent/nodeversion/where.go | 90 ++++++++++++++++ ent/nodeversion_create.go | 65 +++++++++++ ent/nodeversion_update.go | 34 ++++++ ent/runtime.go | 4 + ent/schema/node_version.go | 19 ++++ .../registry_integration_test.go | 55 ++++++++-- node-pack-extract/cloudbuild.yaml | 1 + node-pack-extract/entrypoint.sh | 12 +-- openapi.yml | 2 + server/implementation/registry.go | 17 +-- services/registry/registry_svc.go | 19 +++- 16 files changed, 439 insertions(+), 87 deletions(-) diff --git a/drip/api.gen.go b/drip/api.gen.go index ef501dd..f67c56a 100644 --- a/drip/api.gen.go +++ b/drip/api.gen.go @@ -517,7 +517,8 @@ type ListNodeVersionsParams struct { // CreateComfyNodesJSONBody defines parameters for CreateComfyNodes. type CreateComfyNodesJSONBody struct { - Nodes *map[string]ComfyNode `json:"nodes,omitempty"` + Nodes *map[string]ComfyNode `json:"nodes,omitempty"` + Success *bool `json:"success,omitempty"` } // ValidatePublisherParams defines parameters for ValidatePublisher. @@ -4971,42 +4972,43 @@ var swaggerSpec = []string{ "K5FcbnKwmxfNFiOD4iYtSAZkFi4yV5eTi4sd9tSDpv/bkwXyvqN9m3ivcgnKfvltb4OyEX9bjyy77rS/", "q9G6QX/DFCeSDO/eLF91iYHFNfDrV2s9OZcVKg7/oeUhEJz1bjwokW82ccMB6dNxBxOk5lqU6rJkgKoR", "HDRm7z+Qvz1l3Fft0Y8nXEuKi4ujw2ctin6xqo0yqzs05f4ysVABtnKvlyXcJfJWJJIbPyi3IxYXMoS9", - "GMjHSH15nKSXu3r4ZM2ElXBs1AIC+ovM7UB2yTaZozOZ48Xe/z7cDF5/gwD5GTLpb+boiP5RbiTPMViF", - "PM5gHk4lQmHK6hynERm9Es/ZvY7c90sEdklLG898wpC8hW6MDK8xqC2j+V6zxgrinhFlNYmFR+FogiWJ", - "EWeIKolMXnobfZcX753G82nZ7SFswspLd8ssQu+Zr3KWG5bNs+HqaCUCw0MX79g/3ww79cnyyNZX1nqi", - "RB9V6/l9DRx4wgFeL2lnqpuFhduUsjtNzDXIr003clXSUBsJVTnuLpwPNq/vBK81X81J9LV41KJ8Kse9", - "owRPXCSC4HiBFP5KAteWH+wYPn0uvbwMjKQ4ctNtuw50fR9UIaq9PyQPiniVZoCLyIn/QEixg+6TIQJP", - "KOJqTsQVrTxJ1fFMXBOxLhx0u2VakxBFBs8mXCmmVEptxZT7YN5NgiPdzBs2h8n6CHEbMfjRSj0o8Lr4", - "uzVYYnhHsSnkzPuKHUQUsBE82LckiRctjwfB+s2UQ5LnAfXlcjpP+p0Fc8w+erVrQ2227WMhyd5Dqzxl", - "ls4W9Vbz/pbMqx4DVde48wCOmWdAHhDNHl2bf3DUfgov/aygzW+pcVVBYJ/awSvr9BWNYneCWftFzEvM", - "NlKhMI9SovOte/DJy5yXmFVwuBtb69mo4SAa+YaLh8Hbp5mMcotMkeXhLnIj3yl6aty94vlsy8Po9nf2", - "vrzaMM2rvB96OBdqZ87GU3Cc9n+M60n7L6NcKp52BV52i45KwlK3J+XeyWd4f8lMwewzyD5qc8RsU5Ge", - "lvenGuUdFA3thvkTQ+3HFDEPlBb4nVnzW27ytFwIy7jJCoK1v0PhexSwW8fECo4JLtATyBaoOSh+kMhy", - "8hVoIiMCrjO7MoDeEnVadHtXJRP5VOnkVq/CYfY6piroAen1rEMRdOXt/jbqaoVrIG/jqvVQ/Bctb+d+", - "7siTa8krNV+vmk/zlDXZZoZAX98d1MpcrzqxKy03NqXlxkVRuuZzhOUmfBxYBaIycBuwz2ulKDy/5d6t", - "uR+SzhgJlzyF6r/Q7Oqd5pmtdkrqudY2AmqNNx1PyFUBx2Hzk/BJbQP1+mnjlrNZT9cKWVerMtd6zuJy", - "P9iT4rSdiRX3ofkXxXdYSZc/Wlfbs8f1tbmpxZxI0LMnJOFs1ihN+WC5kI+ftRCYwUo2sztkd8TP1siV", - "HIYLhll7nLPE1Pwuqw9jFnu1gV21qupDWo2xm6G26xUV+/4o+n48isGK09s6Yds6Yds6YQ2fY5i16dX5", - "vK32Ys+q+s9t3C4bHRq7dZT85RwlYDouSYG0VvYB8PT35oNNRuN+mZbNVa0alhQsk//XeoPTrvnpyo1q", - "KmjwQC21hegM0AYdY4ZnJNXrXZosGsC6pxjFHiKe+/ayFa7CQJUacpUsioCq4DGu6R4zJ/yEQrW2brF1", - "8ldDGNNN5z2k6u41/Nkvwe6BGUPY9rXzvXtfltnAv3icmNmE7yxQbE3ScSPuyMgEdgTTwE/zJLGSmUEZ", - "usrDh2DjRZzFeaSQA4g0wGAxDdN+rsfrV4KXsoNZV5G24eDbzozv2B+hCt5hboZsf50yxUVdl9Wff629", - "eB1h1lLyZet0+Q6K+jqURYCzQDbmBnAHC0WnOFLd7yRfQOcD1/euborNeYUVMdNWLc3bVLyGA3w5G18K", - "nLZAuSRQXROnPGdgbn04OzjW5m/s0vVF7mt0XqEPU17RFFMPVSef5NFXoor27hq7pjO6mhNB4AeeK43S", - "tsaxIFAIFPJqGgPBs0TjhM/keBbJMQja/XBRRzVHiiNbhxLpT0aIzEZoJvd3d80cdvS0dnVT+1giZ+Np", - "Anw3NA40VTbRVIsWOWuBmVI1nmM5D26kbfcqCjaHtAUMXZ/2UTTzDIPQLbVSz3b25j1O3fZp8Pvvv/++", - "c3y8c3j4/l//2j8+3j8///enAfrxp73n/9x5vrfzfO/93t4+/PfvZ8F55DH2r/DrBWti7BdIaHxOWNyx", - "Blck1a3jC5+Y1+0vGP0GLVLhNNOQzaIMSv/zRRDDv/DJmMYhTlovvKANVzteaNIakC0cPs4laSFpcDpd", - "zXlZMb4LZoqjOWXm8eWlPphj0/kc+t4MBzyAue/qRbdrXCA0CUOlY6DSgvgCsA8AozCSRKsD5gZMw0Bq", - "jotawBK9fXUORPqjfKbp1OcBLXRqugzbGn4KTToj+GsHT9TNazHETIxtiaUwXK1iWcHQqMXkzW6h5py1", - "UwiAgj7FhWGPg8oWiotovhSs7rQKXEEyHqqVUiktHPwwZyvSVssMoE5yB0so6yjfDVMw1xvLCO43W/r9", - "LPdeVXf14HvKRNc96E/xY7bgEAp2VewtkHlVvFTFQEO0VIV6fcKVvfY4sY/5DfwNML+hU2yKzVwvkOzx", - "yvPWKglHhF5aZu30xWqdqI4SZKKs4XCLguUwg2JwiX70WeczNBU8hek5zectVf/KJ8jUfrcKr+x6ofEt", - "URfSvrx3TzffAD+w2e9+WfUZimbZtjerx4m8JSrwFDkoRzhJ3MNWI2/zfC9Vz6cu33BxF9u6ffFy01+8", - "rFwkWi7YeofYq9iILba9pN5IzSdSxO50+ivJtyyBIGXjkXzweiNaLiwS/YMmv8E9lh/Z1v6/tSS9XeHR", - "Yvo9a6mWLsliY5DrSSV6vrc3um2BU8/r+YCFTocDn+rvoKxNDzVmSaGbx3tCMlSQ9fHK7mxu/Zsquhrh", - "InlKtN6nbKXrlpgrp9qbN0N3r92/z+DfS97W/63Sudf1WR3+xrxEaHTgn/nELib8sH1Kayr9A95m2eGf", - "XuEZ60U0DxDeGO1bAzBIkotksD+YK5XJ/d1dnNEReElHXMwGN59v/j8AAP//KJmqExvDAAA=", + "GMjHSH15nKSXu3r4BE4VDMWgO3utbJZw4NQCov2LtO5A6sk206Mz0+PF3v8+3Axef4Po+RkyuXHm6Ij+", + "UW4kQzJYhTy2YV5VJUJhyursqBE2vRJD2r2O3PdLpHlJaBvPmcKQvIVujID3uVcDV7aM5vtJKSuIe0aU", + "VTMWHoWjCZYkRpwhqiQySett9F3eynda1qdlt4cwGCvP4C0zF703wMpZbliqz4brqpXwDA9dvGP/fDPs", + "VDbLI1tfk+uJEn1Uref3NXDgfQd42qSdqW4WFm7zze40a9cgv7bryFVJQ20kVOW4u3A+2DzNE7zzfDUn", + "0dfixYvyHR33yBK8f5EIguMFUvgrCdxpfrBj+PS59GYzMJLiyE237a7Q9X1Qhaj2OJE8KIJZmtEvIif+", + "6yHFDrpPhgjcpIirORFXtPJe1RLTr0ZPDrrdMq1JiCK9ZxPuG1MqpbZiyn0wjyrBkW7m9ZvDZH2EuI0Y", + "/FCmHhR4XfzdGiwxPLLYFHLm8cUOIgrYCB7sW5LEi5aXhWD9ZsohyfOA+nI5nSf9CIM5Zh+92rWhNtv2", + "sZBk76FVnjKFZ4t6q7mGS+ZVD5Cqa9x5AMfMGyEPiGaPrs0/OGo/hWeAVtDmt9S4qiCw7/DglXX6ikax", + "O8Gs/ZbmJWYbqVCYFyvR+dY9+ORlzkvMKjjcja31VNVwhI18w8XD4O3TzFS5RRrJ8lgYuZGPGD017l7x", + "fLYlaXT7O3tfXm2Y5lXeDz2cC7UzoeMpOE77v9T1pP2XUS4VT7uiMrtFRyWbqduTcu/kM7y/TKdgahqk", + "JrU5YrZ5Sk/L+1MNAQ+KhnbD/Imh9mOKmAfKGfzOrPktN3laLoRl3GQFwdrfofA9CtitY2IFxwQX6Amk", + "EtQcFD9IVITf9qaJjAi4zuxKD3pL1GnR7V2VTORTpZNbPRmH2euYqv7Bza1BV97ub6OuVrgG8jauWizF", + "f+7ydu7njiS6lqRT8/WqyTZPWZNtpg/09d1BIc31She7unNjU3duXFSsa75VWG7Cx4FVICoDtwH7vFaK", + "wvNb7t2a+yHpjJFwPVQoDQzNrhhqntlSqKSeiG0joNZ48PGEXBVwHDY/CZ/UNlCvnzZuOZv1dK2QkrUq", + "c60nNC73gz0pTtuZWHEfmn9RmYeVdPmjdbU9e1xfm5tazIkEPXtCEs5mjbqVD5Yo+fhZC4EZrGQzu0N2", + "R/xsjUTKYbiamLXHOUtMQfCyNDFmsVc42JWyqr6y1Ri7GWq7XsWx74+i78ejGCxHvS0iti0iti0i1vA5", + "hlmbXp3P22rP+ayq/9zG7bLRobFbR8lfzlECpuOSFEhrZR8AT39vPthkNO6Xadlc1aphScEa+n+tBzrt", + "mp+u3KimggYP1FJbiM4AbdAxZnhGUr3epcmiAax7ilHsIeK5by9b4SoMlLAhV8miCKgKHuOa7jFzwk8o", + "VGvrFlsnfzWEMd103kOq7l7Dn/0S7B6YMYRtXzvfu/dlmQ38i8eJmU34zgLF1iQdN+KOjExgRzAN/DRP", + "EiuZGdSoq7yKCDZexFmcRwo5gEgDDFbaMO3nerx+9XkpO5h1VXAbDr7tzPiO/RFK5B3mZsj2pytTXBR9", + "Wf1t2Npz2BFmLfVgtk6X76Dir0NZBDgLZGNuAHewUHSKI9X9iPIFdD5wfe/qpticV1gRM23Vur1NxWs4", + "wJez8aXAaQuUSwKlN3HKcwbm1oezg2Nt/sYuXV/kvkbnVQExtRdNpfVQ6fJJHn0lqmjvLsBrOqOrOREE", + "fuC50ihtCyALAlVCIa+mMRA8SzRO+EyOZ5Ecg6DdD1d8VHOkOLJFKpH+ZITIbIRmcn9318xhR09rVze1", + "jyVyNp4mwHdD40BTZRNNKWmRsxaYKVXjOZbz4Ebadq/cYHNIW93Q9WkfRTPPMAjdUqsDbWdvHuvUbZ8G", + "v//+++87x8c7h4fv//Wv/ePj/fPzf38aoB9/2nv+z53nezvP997v7e3Df/9+FpxHHmP/Cr9ezSbGfvWE", + "xueExR1rcBVU3Tq+8Il5+v6C0W/QIhVOMw3ZLMqg9D9fBDH8C5+MaRzipPWqDNpwteOFJq0B2ari41yS", + "FpIGp9PVnJfl5LtgpjiaU2ZeZl7qgzk2nc+h781wwAOY+65ekbvGBUKTMFQ6BiotiC8A+wAwCiNJtDpg", + "bsA0DKTmuCgULNHbV+dApD/KZ5pOfR7QQqemy7Ct4afQpDOCv3bwRN28FkPMxNjWXwrD1SqWFQyNQk3e", + "7BZqzlk7hQAo6FNcGPY4qGyhuIjmS8HqTqvAFSTjoUIqlbrDwQ9ztiJttcwAiih3sISyyPLdMAVzvbGM", + "4H6zdeHPcu/JdVcsvqdMdN2D/hQ/ZgsOoWBXxd4CmVfFS1UMNERLVajXJ1zZa48T+5jfwN8A8xs6xabY", + "zPUCyR6vdm+tzHBE6KVl1k5frBaR6qhPJsoCD7eoZg4zKAaX6EefdT5DU8FTmJ7TfN5S9a98gkxheKvw", + "yq4XGt8SdSHty3v3dPMN8AOb/e6XVZ+haNZ0e7N6nMhbogLvlINyhJPEPWw18jbP91L1fOryDRd3sa3b", + "Fy83/cXLykWi5YKtd4i9KpHYStxLipHUfCJF7E6nv5J8yxIIUjYeyQcvRqLlwiLRP2jyG9xjbZJh2Ews", + "67/6tc7up/QrSP3i0QnYQj2qXVzx6kTX6Of0T7JCDdi7DV64XVXSYvo9C62WLsliY5DrSSV6vrc3um31", + "U8/r+YBVUIcDn+rvoOZNDzVmSRWcx3tCMlSt9fFq8mxucZwquhrhInlKtN6nbBnslpgrp9qbN0N3r92/", + "z+DfS97W/63Sudf1WR3+xrxEaHTgn/nELib8sH1Kayr9A95m2eGfXlUa60U0DxDeGO1bAzBIkotksD+Y", + "K5XJ/d1dnNEReElHXMwGN59v/j8AAP//O3CizDjDAAA=", } // GetSwagger returns the content of the embedded swagger specification file diff --git a/ent/migrate/schema.go b/ent/migrate/schema.go index c0dc408..56b1b0d 100644 --- a/ent/migrate/schema.go +++ b/ent/migrate/schema.go @@ -187,6 +187,7 @@ var ( {Name: "deprecated", Type: field.TypeBool, Default: false}, {Name: "status", Type: field.TypeEnum, Enums: []string{"active", "banned", "deleted", "pending", "flagged"}, Default: "pending"}, {Name: "status_reason", Type: field.TypeString, Default: "", SchemaType: map[string]string{"postgres": "text"}}, + {Name: "comfy_node_extract_status", Type: field.TypeString, Default: "pending"}, {Name: "node_id", Type: field.TypeString, SchemaType: map[string]string{"postgres": "text"}}, {Name: "node_version_storage_file", Type: field.TypeUUID, Nullable: true}, } @@ -198,13 +199,13 @@ var ( ForeignKeys: []*schema.ForeignKey{ { Symbol: "node_versions_nodes_versions", - Columns: []*schema.Column{NodeVersionsColumns[9]}, + Columns: []*schema.Column{NodeVersionsColumns[10]}, RefColumns: []*schema.Column{NodesColumns[0]}, OnDelete: schema.NoAction, }, { Symbol: "node_versions_storage_files_storage_file", - Columns: []*schema.Column{NodeVersionsColumns[10]}, + Columns: []*schema.Column{NodeVersionsColumns[11]}, RefColumns: []*schema.Column{StorageFilesColumns[0]}, OnDelete: schema.SetNull, }, @@ -213,7 +214,7 @@ var ( { Name: "nodeversion_node_id_version", Unique: true, - Columns: []*schema.Column{NodeVersionsColumns[9], NodeVersionsColumns[3]}, + Columns: []*schema.Column{NodeVersionsColumns[10], NodeVersionsColumns[3]}, }, }, } diff --git a/ent/mutation.go b/ent/mutation.go index 15c9661..f20a3ab 100644 --- a/ent/mutation.go +++ b/ent/mutation.go @@ -6039,29 +6039,30 @@ func (m *NodeReviewMutation) ResetEdge(name string) error { // NodeVersionMutation represents an operation that mutates the NodeVersion nodes in the graph. type NodeVersionMutation struct { config - op Op - typ string - id *uuid.UUID - create_time *time.Time - update_time *time.Time - version *string - changelog *string - pip_dependencies *[]string - appendpip_dependencies []string - deprecated *bool - status *schema.NodeVersionStatus - status_reason *string - clearedFields map[string]struct{} - node *string - clearednode bool - storage_file *uuid.UUID - clearedstorage_file bool - comfy_nodes map[string]struct{} - removedcomfy_nodes map[string]struct{} - clearedcomfy_nodes bool - done bool - oldValue func(context.Context) (*NodeVersion, error) - predicates []predicate.NodeVersion + op Op + typ string + id *uuid.UUID + create_time *time.Time + update_time *time.Time + version *string + changelog *string + pip_dependencies *[]string + appendpip_dependencies []string + deprecated *bool + status *schema.NodeVersionStatus + status_reason *string + comfy_node_extract_status *schema.ComfyNodeExtractStatus + clearedFields map[string]struct{} + node *string + clearednode bool + storage_file *uuid.UUID + clearedstorage_file bool + comfy_nodes map[string]struct{} + removedcomfy_nodes map[string]struct{} + clearedcomfy_nodes bool + done bool + oldValue func(context.Context) (*NodeVersion, error) + predicates []predicate.NodeVersion } var _ ent.Mutation = (*NodeVersionMutation)(nil) @@ -6520,6 +6521,42 @@ func (m *NodeVersionMutation) ResetStatusReason() { m.status_reason = nil } +// SetComfyNodeExtractStatus sets the "comfy_node_extract_status" field. +func (m *NodeVersionMutation) SetComfyNodeExtractStatus(snes schema.ComfyNodeExtractStatus) { + m.comfy_node_extract_status = &snes +} + +// ComfyNodeExtractStatus returns the value of the "comfy_node_extract_status" field in the mutation. +func (m *NodeVersionMutation) ComfyNodeExtractStatus() (r schema.ComfyNodeExtractStatus, exists bool) { + v := m.comfy_node_extract_status + if v == nil { + return + } + return *v, true +} + +// OldComfyNodeExtractStatus returns the old "comfy_node_extract_status" field's value of the NodeVersion entity. +// If the NodeVersion object wasn't provided to the builder, the object is fetched from the database. +// An error is returned if the mutation operation is not UpdateOne, or the database query fails. +func (m *NodeVersionMutation) OldComfyNodeExtractStatus(ctx context.Context) (v schema.ComfyNodeExtractStatus, err error) { + if !m.op.Is(OpUpdateOne) { + return v, errors.New("OldComfyNodeExtractStatus is only allowed on UpdateOne operations") + } + if m.id == nil || m.oldValue == nil { + return v, errors.New("OldComfyNodeExtractStatus requires an ID field in the mutation") + } + oldValue, err := m.oldValue(ctx) + if err != nil { + return v, fmt.Errorf("querying old value for OldComfyNodeExtractStatus: %w", err) + } + return oldValue.ComfyNodeExtractStatus, nil +} + +// ResetComfyNodeExtractStatus resets all changes to the "comfy_node_extract_status" field. +func (m *NodeVersionMutation) ResetComfyNodeExtractStatus() { + m.comfy_node_extract_status = nil +} + // ClearNode clears the "node" edge to the Node entity. func (m *NodeVersionMutation) ClearNode() { m.clearednode = true @@ -6674,7 +6711,7 @@ func (m *NodeVersionMutation) Type() string { // order to get all numeric fields that were incremented/decremented, call // AddedFields(). func (m *NodeVersionMutation) Fields() []string { - fields := make([]string, 0, 9) + fields := make([]string, 0, 10) if m.create_time != nil { fields = append(fields, nodeversion.FieldCreateTime) } @@ -6702,6 +6739,9 @@ func (m *NodeVersionMutation) Fields() []string { if m.status_reason != nil { fields = append(fields, nodeversion.FieldStatusReason) } + if m.comfy_node_extract_status != nil { + fields = append(fields, nodeversion.FieldComfyNodeExtractStatus) + } return fields } @@ -6728,6 +6768,8 @@ func (m *NodeVersionMutation) Field(name string) (ent.Value, bool) { return m.Status() case nodeversion.FieldStatusReason: return m.StatusReason() + case nodeversion.FieldComfyNodeExtractStatus: + return m.ComfyNodeExtractStatus() } return nil, false } @@ -6755,6 +6797,8 @@ func (m *NodeVersionMutation) OldField(ctx context.Context, name string) (ent.Va return m.OldStatus(ctx) case nodeversion.FieldStatusReason: return m.OldStatusReason(ctx) + case nodeversion.FieldComfyNodeExtractStatus: + return m.OldComfyNodeExtractStatus(ctx) } return nil, fmt.Errorf("unknown NodeVersion field %s", name) } @@ -6827,6 +6871,13 @@ func (m *NodeVersionMutation) SetField(name string, value ent.Value) error { } m.SetStatusReason(v) return nil + case nodeversion.FieldComfyNodeExtractStatus: + v, ok := value.(schema.ComfyNodeExtractStatus) + if !ok { + return fmt.Errorf("unexpected type %T for field %s", value, name) + } + m.SetComfyNodeExtractStatus(v) + return nil } return fmt.Errorf("unknown NodeVersion field %s", name) } @@ -6912,6 +6963,9 @@ func (m *NodeVersionMutation) ResetField(name string) error { case nodeversion.FieldStatusReason: m.ResetStatusReason() return nil + case nodeversion.FieldComfyNodeExtractStatus: + m.ResetComfyNodeExtractStatus() + return nil } return fmt.Errorf("unknown NodeVersion field %s", name) } diff --git a/ent/nodeversion.go b/ent/nodeversion.go index 2200106..1786045 100644 --- a/ent/nodeversion.go +++ b/ent/nodeversion.go @@ -40,6 +40,8 @@ type NodeVersion struct { Status schema.NodeVersionStatus `json:"status,omitempty"` // Give a reason for the status change. Eg. 'Banned due to security vulnerability' StatusReason string `json:"status_reason,omitempty"` + // ComfyNodeExtractStatus holds the value of the "comfy_node_extract_status" field. + ComfyNodeExtractStatus schema.ComfyNodeExtractStatus `json:"comfy_node_extract_status,omitempty"` // Edges holds the relations/edges for other nodes in the graph. // The values are being populated by the NodeVersionQuery when eager-loading is set. Edges NodeVersionEdges `json:"edges"` @@ -100,7 +102,7 @@ func (*NodeVersion) scanValues(columns []string) ([]any, error) { values[i] = new([]byte) case nodeversion.FieldDeprecated: values[i] = new(sql.NullBool) - case nodeversion.FieldNodeID, nodeversion.FieldVersion, nodeversion.FieldChangelog, nodeversion.FieldStatus, nodeversion.FieldStatusReason: + case nodeversion.FieldNodeID, nodeversion.FieldVersion, nodeversion.FieldChangelog, nodeversion.FieldStatus, nodeversion.FieldStatusReason, nodeversion.FieldComfyNodeExtractStatus: values[i] = new(sql.NullString) case nodeversion.FieldCreateTime, nodeversion.FieldUpdateTime: values[i] = new(sql.NullTime) @@ -185,6 +187,12 @@ func (nv *NodeVersion) assignValues(columns []string, values []any) error { } else if value.Valid { nv.StatusReason = value.String } + case nodeversion.FieldComfyNodeExtractStatus: + if value, ok := values[i].(*sql.NullString); !ok { + return fmt.Errorf("unexpected type %T for field comfy_node_extract_status", values[i]) + } else if value.Valid { + nv.ComfyNodeExtractStatus = schema.ComfyNodeExtractStatus(value.String) + } case nodeversion.ForeignKeys[0]: if value, ok := values[i].(*sql.NullScanner); !ok { return fmt.Errorf("unexpected type %T for field node_version_storage_file", values[i]) @@ -269,6 +277,9 @@ func (nv *NodeVersion) String() string { builder.WriteString(", ") builder.WriteString("status_reason=") builder.WriteString(nv.StatusReason) + builder.WriteString(", ") + builder.WriteString("comfy_node_extract_status=") + builder.WriteString(fmt.Sprintf("%v", nv.ComfyNodeExtractStatus)) builder.WriteByte(')') return builder.String() } diff --git a/ent/nodeversion/nodeversion.go b/ent/nodeversion/nodeversion.go index 2b3d68c..7f3e63e 100644 --- a/ent/nodeversion/nodeversion.go +++ b/ent/nodeversion/nodeversion.go @@ -35,6 +35,8 @@ const ( FieldStatus = "status" // FieldStatusReason holds the string denoting the status_reason field in the database. FieldStatusReason = "status_reason" + // FieldComfyNodeExtractStatus holds the string denoting the comfy_node_extract_status field in the database. + FieldComfyNodeExtractStatus = "comfy_node_extract_status" // EdgeNode holds the string denoting the node edge name in mutations. EdgeNode = "node" // EdgeStorageFile holds the string denoting the storage_file edge name in mutations. @@ -78,6 +80,7 @@ var Columns = []string{ FieldDeprecated, FieldStatus, FieldStatusReason, + FieldComfyNodeExtractStatus, } // ForeignKeys holds the SQL foreign-keys that are owned by the "node_versions" @@ -112,6 +115,8 @@ var ( DefaultDeprecated bool // DefaultStatusReason holds the default value on creation for the "status_reason" field. DefaultStatusReason string + // DefaultComfyNodeExtractStatus holds the default value on creation for the "comfy_node_extract_status" field. + DefaultComfyNodeExtractStatus schema.ComfyNodeExtractStatus // DefaultID holds the default value on creation for the "id" field. DefaultID func() uuid.UUID ) @@ -176,6 +181,11 @@ func ByStatusReason(opts ...sql.OrderTermOption) OrderOption { return sql.OrderByField(FieldStatusReason, opts...).ToFunc() } +// ByComfyNodeExtractStatus orders the results by the comfy_node_extract_status field. +func ByComfyNodeExtractStatus(opts ...sql.OrderTermOption) OrderOption { + return sql.OrderByField(FieldComfyNodeExtractStatus, opts...).ToFunc() +} + // ByNodeField orders the results by node field. func ByNodeField(field string, opts ...sql.OrderTermOption) OrderOption { return func(s *sql.Selector) { diff --git a/ent/nodeversion/where.go b/ent/nodeversion/where.go index 26dd6f1..2c3457e 100644 --- a/ent/nodeversion/where.go +++ b/ent/nodeversion/where.go @@ -92,6 +92,12 @@ func StatusReason(v string) predicate.NodeVersion { return predicate.NodeVersion(sql.FieldEQ(FieldStatusReason, v)) } +// ComfyNodeExtractStatus applies equality check predicate on the "comfy_node_extract_status" field. It's identical to ComfyNodeExtractStatusEQ. +func ComfyNodeExtractStatus(v schema.ComfyNodeExtractStatus) predicate.NodeVersion { + vc := string(v) + return predicate.NodeVersion(sql.FieldEQ(FieldComfyNodeExtractStatus, vc)) +} + // CreateTimeEQ applies the EQ predicate on the "create_time" field. func CreateTimeEQ(v time.Time) predicate.NodeVersion { return predicate.NodeVersion(sql.FieldEQ(FieldCreateTime, v)) @@ -482,6 +488,90 @@ func StatusReasonContainsFold(v string) predicate.NodeVersion { return predicate.NodeVersion(sql.FieldContainsFold(FieldStatusReason, v)) } +// ComfyNodeExtractStatusEQ applies the EQ predicate on the "comfy_node_extract_status" field. +func ComfyNodeExtractStatusEQ(v schema.ComfyNodeExtractStatus) predicate.NodeVersion { + vc := string(v) + return predicate.NodeVersion(sql.FieldEQ(FieldComfyNodeExtractStatus, vc)) +} + +// ComfyNodeExtractStatusNEQ applies the NEQ predicate on the "comfy_node_extract_status" field. +func ComfyNodeExtractStatusNEQ(v schema.ComfyNodeExtractStatus) predicate.NodeVersion { + vc := string(v) + return predicate.NodeVersion(sql.FieldNEQ(FieldComfyNodeExtractStatus, vc)) +} + +// ComfyNodeExtractStatusIn applies the In predicate on the "comfy_node_extract_status" field. +func ComfyNodeExtractStatusIn(vs ...schema.ComfyNodeExtractStatus) predicate.NodeVersion { + v := make([]any, len(vs)) + for i := range v { + v[i] = string(vs[i]) + } + return predicate.NodeVersion(sql.FieldIn(FieldComfyNodeExtractStatus, v...)) +} + +// ComfyNodeExtractStatusNotIn applies the NotIn predicate on the "comfy_node_extract_status" field. +func ComfyNodeExtractStatusNotIn(vs ...schema.ComfyNodeExtractStatus) predicate.NodeVersion { + v := make([]any, len(vs)) + for i := range v { + v[i] = string(vs[i]) + } + return predicate.NodeVersion(sql.FieldNotIn(FieldComfyNodeExtractStatus, v...)) +} + +// ComfyNodeExtractStatusGT applies the GT predicate on the "comfy_node_extract_status" field. +func ComfyNodeExtractStatusGT(v schema.ComfyNodeExtractStatus) predicate.NodeVersion { + vc := string(v) + return predicate.NodeVersion(sql.FieldGT(FieldComfyNodeExtractStatus, vc)) +} + +// ComfyNodeExtractStatusGTE applies the GTE predicate on the "comfy_node_extract_status" field. +func ComfyNodeExtractStatusGTE(v schema.ComfyNodeExtractStatus) predicate.NodeVersion { + vc := string(v) + return predicate.NodeVersion(sql.FieldGTE(FieldComfyNodeExtractStatus, vc)) +} + +// ComfyNodeExtractStatusLT applies the LT predicate on the "comfy_node_extract_status" field. +func ComfyNodeExtractStatusLT(v schema.ComfyNodeExtractStatus) predicate.NodeVersion { + vc := string(v) + return predicate.NodeVersion(sql.FieldLT(FieldComfyNodeExtractStatus, vc)) +} + +// ComfyNodeExtractStatusLTE applies the LTE predicate on the "comfy_node_extract_status" field. +func ComfyNodeExtractStatusLTE(v schema.ComfyNodeExtractStatus) predicate.NodeVersion { + vc := string(v) + return predicate.NodeVersion(sql.FieldLTE(FieldComfyNodeExtractStatus, vc)) +} + +// ComfyNodeExtractStatusContains applies the Contains predicate on the "comfy_node_extract_status" field. +func ComfyNodeExtractStatusContains(v schema.ComfyNodeExtractStatus) predicate.NodeVersion { + vc := string(v) + return predicate.NodeVersion(sql.FieldContains(FieldComfyNodeExtractStatus, vc)) +} + +// ComfyNodeExtractStatusHasPrefix applies the HasPrefix predicate on the "comfy_node_extract_status" field. +func ComfyNodeExtractStatusHasPrefix(v schema.ComfyNodeExtractStatus) predicate.NodeVersion { + vc := string(v) + return predicate.NodeVersion(sql.FieldHasPrefix(FieldComfyNodeExtractStatus, vc)) +} + +// ComfyNodeExtractStatusHasSuffix applies the HasSuffix predicate on the "comfy_node_extract_status" field. +func ComfyNodeExtractStatusHasSuffix(v schema.ComfyNodeExtractStatus) predicate.NodeVersion { + vc := string(v) + return predicate.NodeVersion(sql.FieldHasSuffix(FieldComfyNodeExtractStatus, vc)) +} + +// ComfyNodeExtractStatusEqualFold applies the EqualFold predicate on the "comfy_node_extract_status" field. +func ComfyNodeExtractStatusEqualFold(v schema.ComfyNodeExtractStatus) predicate.NodeVersion { + vc := string(v) + return predicate.NodeVersion(sql.FieldEqualFold(FieldComfyNodeExtractStatus, vc)) +} + +// ComfyNodeExtractStatusContainsFold applies the ContainsFold predicate on the "comfy_node_extract_status" field. +func ComfyNodeExtractStatusContainsFold(v schema.ComfyNodeExtractStatus) predicate.NodeVersion { + vc := string(v) + return predicate.NodeVersion(sql.FieldContainsFold(FieldComfyNodeExtractStatus, vc)) +} + // HasNode applies the HasEdge predicate on the "node" edge. func HasNode() predicate.NodeVersion { return predicate.NodeVersion(func(s *sql.Selector) { diff --git a/ent/nodeversion_create.go b/ent/nodeversion_create.go index 1fac388..1490c17 100644 --- a/ent/nodeversion_create.go +++ b/ent/nodeversion_create.go @@ -130,6 +130,20 @@ func (nvc *NodeVersionCreate) SetNillableStatusReason(s *string) *NodeVersionCre return nvc } +// SetComfyNodeExtractStatus sets the "comfy_node_extract_status" field. +func (nvc *NodeVersionCreate) SetComfyNodeExtractStatus(snes schema.ComfyNodeExtractStatus) *NodeVersionCreate { + nvc.mutation.SetComfyNodeExtractStatus(snes) + return nvc +} + +// SetNillableComfyNodeExtractStatus sets the "comfy_node_extract_status" field if the given value is not nil. +func (nvc *NodeVersionCreate) SetNillableComfyNodeExtractStatus(snes *schema.ComfyNodeExtractStatus) *NodeVersionCreate { + if snes != nil { + nvc.SetComfyNodeExtractStatus(*snes) + } + return nvc +} + // SetID sets the "id" field. func (nvc *NodeVersionCreate) SetID(u uuid.UUID) *NodeVersionCreate { nvc.mutation.SetID(u) @@ -238,6 +252,10 @@ func (nvc *NodeVersionCreate) defaults() { v := nodeversion.DefaultStatusReason nvc.mutation.SetStatusReason(v) } + if _, ok := nvc.mutation.ComfyNodeExtractStatus(); !ok { + v := nodeversion.DefaultComfyNodeExtractStatus + nvc.mutation.SetComfyNodeExtractStatus(v) + } if _, ok := nvc.mutation.ID(); !ok { v := nodeversion.DefaultID() nvc.mutation.SetID(v) @@ -275,6 +293,9 @@ func (nvc *NodeVersionCreate) check() error { if _, ok := nvc.mutation.StatusReason(); !ok { return &ValidationError{Name: "status_reason", err: errors.New(`ent: missing required field "NodeVersion.status_reason"`)} } + if _, ok := nvc.mutation.ComfyNodeExtractStatus(); !ok { + return &ValidationError{Name: "comfy_node_extract_status", err: errors.New(`ent: missing required field "NodeVersion.comfy_node_extract_status"`)} + } if _, ok := nvc.mutation.NodeID(); !ok { return &ValidationError{Name: "node", err: errors.New(`ent: missing required edge "NodeVersion.node"`)} } @@ -346,6 +367,10 @@ func (nvc *NodeVersionCreate) createSpec() (*NodeVersion, *sqlgraph.CreateSpec) _spec.SetField(nodeversion.FieldStatusReason, field.TypeString, value) _node.StatusReason = value } + if value, ok := nvc.mutation.ComfyNodeExtractStatus(); ok { + _spec.SetField(nodeversion.FieldComfyNodeExtractStatus, field.TypeString, value) + _node.ComfyNodeExtractStatus = value + } if nodes := nvc.mutation.NodeIDs(); len(nodes) > 0 { edge := &sqlgraph.EdgeSpec{ Rel: sqlgraph.M2O, @@ -550,6 +575,18 @@ func (u *NodeVersionUpsert) UpdateStatusReason() *NodeVersionUpsert { return u } +// SetComfyNodeExtractStatus sets the "comfy_node_extract_status" field. +func (u *NodeVersionUpsert) SetComfyNodeExtractStatus(v schema.ComfyNodeExtractStatus) *NodeVersionUpsert { + u.Set(nodeversion.FieldComfyNodeExtractStatus, v) + return u +} + +// UpdateComfyNodeExtractStatus sets the "comfy_node_extract_status" field to the value that was provided on create. +func (u *NodeVersionUpsert) UpdateComfyNodeExtractStatus() *NodeVersionUpsert { + u.SetExcluded(nodeversion.FieldComfyNodeExtractStatus) + return u +} + // UpdateNewValues updates the mutable fields using the new values that were set on create except the ID field. // Using this option is equivalent to using: // @@ -720,6 +757,20 @@ func (u *NodeVersionUpsertOne) UpdateStatusReason() *NodeVersionUpsertOne { }) } +// SetComfyNodeExtractStatus sets the "comfy_node_extract_status" field. +func (u *NodeVersionUpsertOne) SetComfyNodeExtractStatus(v schema.ComfyNodeExtractStatus) *NodeVersionUpsertOne { + return u.Update(func(s *NodeVersionUpsert) { + s.SetComfyNodeExtractStatus(v) + }) +} + +// UpdateComfyNodeExtractStatus sets the "comfy_node_extract_status" field to the value that was provided on create. +func (u *NodeVersionUpsertOne) UpdateComfyNodeExtractStatus() *NodeVersionUpsertOne { + return u.Update(func(s *NodeVersionUpsert) { + s.UpdateComfyNodeExtractStatus() + }) +} + // Exec executes the query. func (u *NodeVersionUpsertOne) Exec(ctx context.Context) error { if len(u.create.conflict) == 0 { @@ -1057,6 +1108,20 @@ func (u *NodeVersionUpsertBulk) UpdateStatusReason() *NodeVersionUpsertBulk { }) } +// SetComfyNodeExtractStatus sets the "comfy_node_extract_status" field. +func (u *NodeVersionUpsertBulk) SetComfyNodeExtractStatus(v schema.ComfyNodeExtractStatus) *NodeVersionUpsertBulk { + return u.Update(func(s *NodeVersionUpsert) { + s.SetComfyNodeExtractStatus(v) + }) +} + +// UpdateComfyNodeExtractStatus sets the "comfy_node_extract_status" field to the value that was provided on create. +func (u *NodeVersionUpsertBulk) UpdateComfyNodeExtractStatus() *NodeVersionUpsertBulk { + return u.Update(func(s *NodeVersionUpsert) { + s.UpdateComfyNodeExtractStatus() + }) +} + // Exec executes the query. func (u *NodeVersionUpsertBulk) Exec(ctx context.Context) error { if u.create.err != nil { diff --git a/ent/nodeversion_update.go b/ent/nodeversion_update.go index 0160c1a..586dbb1 100644 --- a/ent/nodeversion_update.go +++ b/ent/nodeversion_update.go @@ -143,6 +143,20 @@ func (nvu *NodeVersionUpdate) SetNillableStatusReason(s *string) *NodeVersionUpd return nvu } +// SetComfyNodeExtractStatus sets the "comfy_node_extract_status" field. +func (nvu *NodeVersionUpdate) SetComfyNodeExtractStatus(snes schema.ComfyNodeExtractStatus) *NodeVersionUpdate { + nvu.mutation.SetComfyNodeExtractStatus(snes) + return nvu +} + +// SetNillableComfyNodeExtractStatus sets the "comfy_node_extract_status" field if the given value is not nil. +func (nvu *NodeVersionUpdate) SetNillableComfyNodeExtractStatus(snes *schema.ComfyNodeExtractStatus) *NodeVersionUpdate { + if snes != nil { + nvu.SetComfyNodeExtractStatus(*snes) + } + return nvu +} + // SetNode sets the "node" edge to the Node entity. func (nvu *NodeVersionUpdate) SetNode(n *Node) *NodeVersionUpdate { return nvu.SetNodeID(n.ID) @@ -316,6 +330,9 @@ func (nvu *NodeVersionUpdate) sqlSave(ctx context.Context) (n int, err error) { if value, ok := nvu.mutation.StatusReason(); ok { _spec.SetField(nodeversion.FieldStatusReason, field.TypeString, value) } + if value, ok := nvu.mutation.ComfyNodeExtractStatus(); ok { + _spec.SetField(nodeversion.FieldComfyNodeExtractStatus, field.TypeString, value) + } if nvu.mutation.NodeCleared() { edge := &sqlgraph.EdgeSpec{ Rel: sqlgraph.M2O, @@ -549,6 +566,20 @@ func (nvuo *NodeVersionUpdateOne) SetNillableStatusReason(s *string) *NodeVersio return nvuo } +// SetComfyNodeExtractStatus sets the "comfy_node_extract_status" field. +func (nvuo *NodeVersionUpdateOne) SetComfyNodeExtractStatus(snes schema.ComfyNodeExtractStatus) *NodeVersionUpdateOne { + nvuo.mutation.SetComfyNodeExtractStatus(snes) + return nvuo +} + +// SetNillableComfyNodeExtractStatus sets the "comfy_node_extract_status" field if the given value is not nil. +func (nvuo *NodeVersionUpdateOne) SetNillableComfyNodeExtractStatus(snes *schema.ComfyNodeExtractStatus) *NodeVersionUpdateOne { + if snes != nil { + nvuo.SetComfyNodeExtractStatus(*snes) + } + return nvuo +} + // SetNode sets the "node" edge to the Node entity. func (nvuo *NodeVersionUpdateOne) SetNode(n *Node) *NodeVersionUpdateOne { return nvuo.SetNodeID(n.ID) @@ -752,6 +783,9 @@ func (nvuo *NodeVersionUpdateOne) sqlSave(ctx context.Context) (_node *NodeVersi if value, ok := nvuo.mutation.StatusReason(); ok { _spec.SetField(nodeversion.FieldStatusReason, field.TypeString, value) } + if value, ok := nvuo.mutation.ComfyNodeExtractStatus(); ok { + _spec.SetField(nodeversion.FieldComfyNodeExtractStatus, field.TypeString, value) + } if nvuo.mutation.NodeCleared() { edge := &sqlgraph.EdgeSpec{ Rel: sqlgraph.M2O, diff --git a/ent/runtime.go b/ent/runtime.go index 5226394..e867e17 100644 --- a/ent/runtime.go +++ b/ent/runtime.go @@ -160,6 +160,10 @@ func init() { nodeversionDescStatusReason := nodeversionFields[7].Descriptor() // nodeversion.DefaultStatusReason holds the default value on creation for the status_reason field. nodeversion.DefaultStatusReason = nodeversionDescStatusReason.Default.(string) + // nodeversionDescComfyNodeExtractStatus is the schema descriptor for comfy_node_extract_status field. + nodeversionDescComfyNodeExtractStatus := nodeversionFields[8].Descriptor() + // nodeversion.DefaultComfyNodeExtractStatus holds the default value on creation for the comfy_node_extract_status field. + nodeversion.DefaultComfyNodeExtractStatus = schema.ComfyNodeExtractStatus(nodeversionDescComfyNodeExtractStatus.Default.(string)) // nodeversionDescID is the schema descriptor for id field. nodeversionDescID := nodeversionFields[0].Descriptor() // nodeversion.DefaultID holds the default value on creation for the id field. diff --git a/ent/schema/node_version.go b/ent/schema/node_version.go index cc8cba1..7229373 100644 --- a/ent/schema/node_version.go +++ b/ent/schema/node_version.go @@ -38,6 +38,9 @@ func (NodeVersion) Fields() []ent.Field { field.String("status_reason").SchemaType(map[string]string{ dialect.Postgres: "text", }).Default("").Comment("Give a reason for the status change. Eg. 'Banned due to security vulnerability'"), + field.String("comfy_node_extract_status"). + GoType(ComfyNodeExtractStatus("")). + Default(string(ComfyNodeExtractStatusPending)), } } @@ -80,3 +83,19 @@ func (NodeVersionStatus) Values() (types []string) { string(NodeVersionStatusFlagged), } } + +type ComfyNodeExtractStatus string + +func (ComfyNodeExtractStatus) Values() (types []string) { + return []string{ + string(ComfyNodeExtractStatusPending), + string(ComfyNodeExtractStatusFailed), + string(ComfyNodeExtractStatusSuccess), + } +} + +const ( + ComfyNodeExtractStatusPending = "pending" + ComfyNodeExtractStatusFailed = "failed" + ComfyNodeExtractStatusSuccess = "success" +) diff --git a/integration-tests/registry_integration_test.go b/integration-tests/registry_integration_test.go index 860c458..59f554a 100644 --- a/integration-tests/registry_integration_test.go +++ b/integration-tests/registry_integration_test.go @@ -838,7 +838,8 @@ func TestRegistryComfyNode(t *testing.T) { require.NoError(t, err, "should not return error") } nodeVersion := nodeVersions[len(nodeVersions)-1] - backfilledNodeVersions := nodeVersions[:len(nodeVersions)-1] + nodeVersionExtractionFailed := nodeVersions[len(nodeVersions)-2] + backfilledNodeVersions := nodeVersions[:len(nodeVersions)-2] t.Run("NoComfyNode", func(t *testing.T) { res, err := withMiddleware(authz, impl.GetNodeVersion)(ctx, drip.GetNodeVersionRequestObject{ @@ -877,14 +878,50 @@ func TestRegistryComfyNode(t *testing.T) { }} // create comfy nodes - body := drip.CreateComfyNodesJSONRequestBody(comfyNodes) - res, err := withMiddleware(authz, impl.CreateComfyNodes)(ctx, drip.CreateComfyNodesRequestObject{ - NodeId: *node.Id, - Version: *nodeVersion.Version, - Body: &body, - }) - require.NoError(t, err) - require.IsType(t, drip.CreateComfyNodes204Response{}, res) + { + body := drip.CreateComfyNodesJSONRequestBody(comfyNodes) + res, err := withMiddleware(authz, impl.CreateComfyNodes)(ctx, drip.CreateComfyNodesRequestObject{ + NodeId: *node.Id, + Version: *nodeVersion.Version, + Body: &body, + }) + require.NoError(t, err) + require.IsType(t, drip.CreateComfyNodes204Response{}, res) + } + + // mark comfy nodes extraction as failed + { + res, err := withMiddleware(authz, impl.CreateComfyNodes)(ctx, drip.CreateComfyNodesRequestObject{ + NodeId: *node.Id, + Version: *nodeVersionExtractionFailed.Version, + Body: &drip.CreateComfyNodesJSONRequestBody{Success: proto.Bool(false)}, + }) + require.NoError(t, err) + require.IsType(t, drip.CreateComfyNodes204Response{}, res) + } + + t.Run("AssertAlgolia", func(t *testing.T) { + indexed := impl.mockAlgolia.LastIndexedNodes + require.Len(t, impl.mockAlgolia.LastIndexedNodes, 1) + + node, err := client.Node.Get(ctx, *node.Id) + require.NoError(t, err) + nodeVersion, err := client.NodeVersion.Query().Where(nodeversion.Version(*nodeVersion.Version)).WithComfyNodes().Only(ctx) + require.NoError(t, err) + node.Edges.Versions = append(node.Edges.Versions, nodeVersion) + + assert.Equal(t, node.ID, indexed[0].ID) + assert.Equal(t, node.Edges.Versions[0].ID, indexed[0].Edges.Versions[0].ID) + indexedComfyNodes := drip.CreateComfyNodesJSONRequestBody{ + Nodes: &map[string]drip.ComfyNode{}, + } + for _, node := range indexed[0].Edges.Versions[0].Edges.ComfyNodes { + cn := *(mapper.DBComfyNodeToApiComfyNode(node)) + cn.ComfyNodeId = nil + (*indexedComfyNodes.Nodes)[node.ID] = cn + } + assert.Equal(t, comfyNodes, indexedComfyNodes) + }) t.Run("AssertAlgolia", func(t *testing.T) { indexed := impl.mockAlgolia.LastIndexedNodes diff --git a/node-pack-extract/cloudbuild.yaml b/node-pack-extract/cloudbuild.yaml index c0137b6..72f4fbe 100644 --- a/node-pack-extract/cloudbuild.yaml +++ b/node-pack-extract/cloudbuild.yaml @@ -28,6 +28,7 @@ steps: args: - -c - gcloud auth print-identity-token --audiences="$_REGISTRY_BACKEND_URL" | tee /workspace/token + waitFor: ["-"] - name: "curlimages/curl" entrypoint: "sh" diff --git a/node-pack-extract/entrypoint.sh b/node-pack-extract/entrypoint.sh index 7f4c6f4..34d9bbd 100644 --- a/node-pack-extract/entrypoint.sh +++ b/node-pack-extract/entrypoint.sh @@ -8,6 +8,7 @@ init.sh & OUTPUTFILE=${1:-"/tmp/output.json"} echo -n > "$OUTPUTFILE" until cat "$OUTPUTFILE" | grep ''; do + sleep 1 curl -sf localhost:8188/object_info | jq -c ' to_entries | @@ -24,11 +25,10 @@ until cat "$OUTPUTFILE" | grep ''; do output_is_list : .output_is_list, } ) | - if length > 0 then {nodes: from_entries} else "" end' | + if length > 0 then + {success: true, nodes: from_entries} + else + {success: false} + end' | tee "$OUTPUTFILE" - - sleep 1 done - -# make sure its json or we fail -grep '{' "$OUTPUTFILE" diff --git a/openapi.yml b/openapi.yml index 02568ac..31752d3 100644 --- a/openapi.yml +++ b/openapi.yml @@ -1788,6 +1788,8 @@ paths: schema: type: object properties: + success: + type: boolean nodes: additionalProperties: $ref: '#/components/schemas/ComfyNode' diff --git a/server/implementation/registry.go b/server/implementation/registry.go index a1bd34e..a3aaa25 100644 --- a/server/implementation/registry.go +++ b/server/implementation/registry.go @@ -1008,20 +1008,25 @@ func (s *DripStrictServerImplementation) ReindexNodes(ctx context.Context, reque return drip.ReindexNodes200Response{}, nil } -// CreateComfyNodes bulk-creates comfy-nodes for a node version +// CreateComfyNodes bulk-stores comfy-nodes extraction result for a node version func (impl *DripStrictServerImplementation) CreateComfyNodes(ctx context.Context, request drip.CreateComfyNodesRequestObject) (res drip.CreateComfyNodesResponseObject, err error) { - err = impl.RegistryService.CreateComfyNodes(ctx, impl.Client, request.NodeId, request.Version, *request.Body.Nodes) + if request.Body.Success != nil && !*request.Body.Success { + err = impl.RegistryService.MarKComfyNodeExtractionFailed(ctx, impl.Client, request.NodeId, request.Version) + } else { + err = impl.RegistryService.CreateComfyNodes(ctx, impl.Client, request.NodeId, request.Version, *request.Body.Nodes) + } + if ent.IsNotFound(err) { log.Ctx(ctx).Error().Msgf("Node or node version not found w/ err: %v", err) return drip.CreateComfyNodes404JSONResponse{Message: "Node or node version not found", Error: err.Error()}, nil } if errors.Is(err, drip_services.ErrComfyNodesAlreadyExist) { - log.Ctx(ctx).Error().Msgf("Comfy nodes for %s %s exist", request.NodeId, request.Version) - return drip.CreateComfyNodes409JSONResponse{Message: "Comfy nodes already exist", Error: err.Error()}, nil + log.Ctx(ctx).Error().Msgf("Comfy nodes extraction result for %s %s already set", request.NodeId, request.Version) + return drip.CreateComfyNodes409JSONResponse{Message: "Comfy nodes extraction result already set", Error: err.Error()}, nil } if err != nil { - log.Ctx(ctx).Error().Msgf("Failed to create comfy nodes w/ err: %v", err) - return drip.CreateComfyNodes500JSONResponse{Message: "Failed to create comfy nodes", Error: err.Error()}, nil + log.Ctx(ctx).Error().Msgf("Failed to store comfy nodes extraction w/ err: %v", err) + return drip.CreateComfyNodes500JSONResponse{Message: "Failed to store comfy nodes extraction", Error: err.Error()}, nil } log.Ctx(ctx).Info().Msgf("CreateComfyNodes successful") diff --git a/services/registry/registry_svc.go b/services/registry/registry_svc.go index 1f01915..1474098 100644 --- a/services/registry/registry_svc.go +++ b/services/registry/registry_svc.go @@ -570,6 +570,17 @@ func (s *RegistryService) GetLatestNodeVersion(ctx context.Context, client *ent. var ErrComfyNodesAlreadyExist = errors.New("comfy nodes already exist") +func (s *RegistryService) MarKComfyNodeExtractionFailed(ctx context.Context, client *ent.Client, nodeID string, nodeVersion string) error { + return client.NodeVersion. + Update(). + Where( + nodeversion.NodeIDEQ(nodeID), + nodeversion.VersionEQ(nodeVersion), + ). + SetComfyNodeExtractStatus(schema.ComfyNodeExtractStatusFailed). + Exec(ctx) +} + func (s *RegistryService) CreateComfyNodes(ctx context.Context, client *ent.Client, nodeID string, nodeVersion string, comfyNodes map[string]drip.ComfyNode) (err error) { return db.WithTx(ctx, client, func(tx *ent.Tx) error { nv, err := tx.NodeVersion.Query(). @@ -629,6 +640,11 @@ func (s *RegistryService) CreateComfyNodes(ctx context.Context, client *ent.Clie return fmt.Errorf("failed to update comfy nodes: %w", err) } + err = nv.Update().SetComfyNodeExtractStatus(schema.ComfyNodeExtractStatusSuccess).Exec(ctx) + if err != nil { + return fmt.Errorf("failed to update comfy nodes extraction status: %w", err) + } + if _, err := s.indexNodeWithLatestVersion(ctx, tx.Client(), nodeID); err != nil { return fmt.Errorf("failed to update node index") } @@ -642,6 +658,7 @@ func (s *RegistryService) GetComfyNode(ctx context.Context, client *ent.Client, nv, err := client.NodeVersion.Query(). Where(nodeversion.VersionEQ(nodeVersion)). Where(nodeversion.NodeIDEQ(nodeID)). + Where(nodeversion.ComfyNodeExtractStatusEQ(schema.ComfyNodeExtractStatusSuccess)). WithComfyNodes(func(cnq *ent.ComfyNodeQuery) { cnq.Where(comfynode.IDEQ(comfyNodeID)) }). @@ -656,7 +673,7 @@ func (s *RegistryService) TriggerComfyNodesBackfill(ctx context.Context, client q := client.NodeVersion. Query(). WithStorageFile(). - Where(nodeversion.Not(nodeversion.HasComfyNodes())) + Where(nodeversion.ComfyNodeExtractStatusEQ(schema.ComfyNodeExtractStatusPending)) if max != nil { q.Limit(*max) }