Skip to content

Commit aeeacb0

Browse files
naveenku-jfrognaveenku-jfrog
committed
Optimized AQL and funcs
1 parent 8d510cc commit aeeacb0

File tree

3 files changed

+45
-90
lines changed

3 files changed

+45
-90
lines changed

‎artifactory/commands/huggingface/huggingFaceDownload.go‎

Lines changed: 2 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ import (
77
"io"
88
"os"
99
"os/exec"
10-
"sort"
1110

1211
"github.com/jfrog/build-info-go/entities"
1312
"github.com/jfrog/jfrog-cli-core/v2/artifactory/utils"
@@ -137,12 +136,10 @@ func (hfd *HuggingFaceDownload) GetDependencies() ([]entities.Dependency, error)
137136
}
138137
repoTypePath := hfd.repoType + "s"
139138
revisionPattern := hfd.revision
140-
var multipleDirsInSearchResults = false
141139
if !HasTimestamp(hfd.revision) {
142140
revisionPattern = hfd.revision + "_*"
143-
multipleDirsInSearchResults = true
144141
}
145-
aqlQuery := fmt.Sprintf(`items.find({"repo":"%s","path":{"$match":"%s/%s/%s/*"}}).include("repo","path","name","actual_sha1","actual_md5","sha256","type")`,
142+
aqlQuery := fmt.Sprintf(`items.find({"repo":"%s","path":{"$match":"%s/%s/%s/*"}}).include("repo","path","name","actual_sha1","actual_md5","sha256","type").sort({"$desc":["path"]})`,
146143
repoKey,
147144
repoTypePath,
148145
hfd.repoId,
@@ -155,17 +152,8 @@ func (hfd *HuggingFaceDownload) GetDependencies() ([]entities.Dependency, error)
155152
if len(results) == 0 {
156153
return nil, nil
157154
}
158-
if multipleDirsInSearchResults {
159-
sort.Slice(results, func(i, j int) bool {
160-
return results[i].Path > results[j].Path
161-
})
162-
}
163-
var latestCreatedDir string
164155
var dependencies []entities.Dependency
165-
for index, resultItem := range results {
166-
if index == 0 {
167-
latestCreatedDir = resultItem.Path
168-
}
156+
for _, resultItem := range results {
169157
dependencies = append(dependencies, entities.Dependency{
170158
Id: resultItem.Name,
171159
Type: resultItem.Type,
@@ -176,9 +164,6 @@ func (hfd *HuggingFaceDownload) GetDependencies() ([]entities.Dependency, error)
176164
Sha256: resultItem.Sha256,
177165
},
178166
})
179-
if latestCreatedDir != resultItem.Path {
180-
break
181-
}
182167
}
183168
return dependencies, nil
184169
}

‎artifactory/commands/huggingface/huggingFaceUpload.go‎

Lines changed: 16 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -4,19 +4,15 @@ import (
44
"bytes"
55
"encoding/json"
66
"fmt"
7-
"github.com/jfrog/build-info-go/entities"
8-
"github.com/jfrog/jfrog-cli-core/v2/artifactory/utils"
9-
buildUtils "github.com/jfrog/jfrog-cli-core/v2/common/build"
10-
"github.com/jfrog/jfrog-client-go/artifactory/services"
11-
servicesUtils "github.com/jfrog/jfrog-client-go/artifactory/services/utils"
12-
"github.com/jfrog/jfrog-client-go/utils/io/content"
137
"io"
148
"os"
159
"os/exec"
16-
"sort"
1710
"strconv"
1811
"time"
1912

13+
"github.com/jfrog/build-info-go/entities"
14+
"github.com/jfrog/jfrog-cli-core/v2/artifactory/utils"
15+
buildUtils "github.com/jfrog/jfrog-cli-core/v2/common/build"
2016
"github.com/jfrog/jfrog-cli-core/v2/utils/config"
2117
"github.com/jfrog/jfrog-client-go/utils/errorutils"
2218
"github.com/jfrog/jfrog-client-go/utils/log"
@@ -141,50 +137,25 @@ func (hfu *HuggingFaceUpload) GetArtifacts(buildProperties string) ([]entities.A
141137
}
142138
repoTypePath := hfu.repoType + "s"
143139
revisionPattern := hfu.revision
144-
var multipleDirsInSearchResults = false
145140
if !HasTimestamp(hfu.revision) {
146141
revisionPattern = hfu.revision + "_*"
147-
multipleDirsInSearchResults = true
148142
}
149-
aqlQuery := fmt.Sprintf(`{"repo": "%s", "path": {"$match": "%s/%s/%s/*"}}`,
143+
aqlQuery := fmt.Sprintf(`items.find({"repo":"%s","path":{"$match":"%s/%s/%s/*"}}).include("repo","path","name","actual_sha1","actual_md5","sha256","type").sort({"$desc":["path"]})`,
150144
repoKey,
151145
repoTypePath,
152146
hfu.repoId,
153147
revisionPattern,
154148
)
155-
searchParams := services.SearchParams{
156-
CommonParams: &servicesUtils.CommonParams{
157-
Aql: servicesUtils.Aql{ItemsFind: aqlQuery},
158-
},
159-
}
160-
reader, err := serviceManager.SearchFiles(searchParams)
149+
results, err := executeAqlQuery(serviceManager, aqlQuery)
161150
if err != nil {
162151
return nil, fmt.Errorf("failed to search for HuggingFace artifacts: %w", err)
163152
}
164-
defer func(reader *content.ContentReader) {
165-
err := reader.Close()
166-
if err != nil {
167-
log.Error(err)
168-
}
169-
}(reader)
170-
var results []servicesUtils.ResultItem
171-
for item := new(servicesUtils.ResultItem); reader.NextRecord(item) == nil; item = new(servicesUtils.ResultItem) {
172-
results = append(results, *item)
173-
}
174153
if len(results) == 0 {
175154
return nil, nil
176155
}
177-
if multipleDirsInSearchResults {
178-
sort.Slice(results, func(i, j int) bool {
179-
return results[i].Path > results[j].Path
180-
})
181-
}
182-
var latestCreatedDir string
156+
latestCreatedDir := results[0].Path
183157
var artifacts []entities.Artifact
184-
for index, resultItem := range results {
185-
if index == 0 {
186-
latestCreatedDir = resultItem.Path
187-
}
158+
for _, resultItem := range results {
188159
artifacts = append(artifacts, entities.Artifact{
189160
Name: resultItem.Name,
190161
Type: resultItem.Type,
@@ -194,15 +165,18 @@ func (hfu *HuggingFaceUpload) GetArtifacts(buildProperties string) ([]entities.A
194165
Sha256: resultItem.Sha256,
195166
},
196167
})
197-
if latestCreatedDir != resultItem.Path {
198-
break
199-
}
200168
}
201-
err = updateReaderContents(reader, repoKey, latestCreatedDir, "")
169+
// Create content reader for the folder to set build properties
170+
reader, err := createContentReader(repoKey, latestCreatedDir, "", "folder")
202171
if err != nil {
203-
return nil, err
172+
log.Warn("Failed to create content reader: ", err)
173+
return artifacts, nil
204174
}
205-
reader.Reset()
175+
defer func() {
176+
if closeErr := reader.Close(); closeErr != nil {
177+
log.Error(closeErr)
178+
}
179+
}()
206180
addBuildPropertiesOnArtifacts(serviceManager, reader, buildProperties)
207181
return artifacts, nil
208182
}

‎artifactory/commands/huggingface/utils.go‎

Lines changed: 27 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -139,37 +139,6 @@ func GetRepoKeyFromHFEndpoint() (string, error) {
139139
return repoKey, nil
140140
}
141141

142-
// updateReaderContents updates the reader contents by writing the specified JSON value to all file paths
143-
func updateReaderContents(reader *content.ContentReader, repo, path, name string) error {
144-
if reader == nil {
145-
return fmt.Errorf("reader is nil")
146-
}
147-
jsonData := map[string]interface{}{
148-
"results": []map[string]interface{}{
149-
{
150-
"repo": repo,
151-
"path": path,
152-
"name": name,
153-
"type": "folder",
154-
},
155-
},
156-
}
157-
jsonBytes, err := json.Marshal(jsonData)
158-
if err != nil {
159-
return fmt.Errorf("failed to marshal JSON: %w", err)
160-
}
161-
filesPaths := reader.GetFilesPaths()
162-
for _, filePath := range filesPaths {
163-
err := os.WriteFile(filePath, jsonBytes, 0644)
164-
if err != nil {
165-
log.Warn(fmt.Sprintf("Failed to write JSON to file %s: %s", filePath, err))
166-
continue
167-
}
168-
log.Debug(fmt.Sprintf("Successfully updated file %s with JSON content", filePath))
169-
}
170-
return nil
171-
}
172-
173142
// executeAqlQuery executes an AQL query and parses the JSON response
174143
func executeAqlQuery(serviceManager artifactory.ArtifactoryServicesManager, aqlQuery string) ([]servicesUtils.ResultItem, error) {
175144
reader, err := serviceManager.Aql(aqlQuery)
@@ -195,6 +164,33 @@ func executeAqlQuery(serviceManager artifactory.ArtifactoryServicesManager, aqlQ
195164
return parsedResult.Results, nil
196165
}
197166

167+
type aqlResult struct {
168+
Results []aqlItem `json:"results"`
169+
}
170+
171+
type aqlItem struct {
172+
Repo string `json:"repo"`
173+
Path string `json:"path"`
174+
Name string `json:"name"`
175+
Type string `json:"type"`
176+
}
177+
178+
// createContentReader creates a ContentReader from the provided parameters
179+
func createContentReader(repo, path, name, itemType string) (*content.ContentReader, error) {
180+
tempFile, err := os.CreateTemp("", "aql-results-*.json")
181+
if err != nil {
182+
return nil, fmt.Errorf("failed to create temp file: %w", err)
183+
}
184+
filePath := tempFile.Name()
185+
err = json.NewEncoder(tempFile).Encode(aqlResult{
186+
Results: []aqlItem{{Repo: repo, Path: path, Name: name, Type: itemType}},
187+
})
188+
if err != nil {
189+
return nil, fmt.Errorf("failed to write to temp file: %w", err)
190+
}
191+
return content.NewContentReader(filePath, content.DefaultKey), nil
192+
}
193+
198194
func addBuildPropertiesOnArtifacts(serviceManager artifactory.ArtifactoryServicesManager, reader *content.ContentReader, buildProps string) {
199195
propsParams := services.PropsParams{
200196
Reader: reader,

0 commit comments

Comments
 (0)