Merge branch 'jdp-missing-files' into add-jdp-datasets-search-param

Author: jeff-cohere

databases/jdp/database.go

@@ -75,8 +75,8 @@ type Config struct {
 }
 
 type StagingRequest struct {
-	// JDP staging request ID
-	Id int
+	// JDP staging request IDs (batched because the endpoint has a maximum limit)
+	Ids []int
 	// time of staging request (for purging)
 	Time time.Time
 }
@@ -173,20 +173,13 @@ func (db *Database) Search(orcid string, params databases.SearchParameters) (dat
 	}, err
 }
 
-func (db *Database) Descriptors(orcid string, fileIds []string) ([]map[string]any, error) {
-	// strip the "JDP:" prefix from our files and create a mapping from IDs to
-	// their original order so we can hand back metadata accordingly
+func (db *Database) fetchDescriptors(orcid string, fileIds []string, batchSize int) ([]map[string]any, error) {
+	// strip the "JDP:" prefix from our files
 	strippedFileIds := make([]string, len(fileIds))
-	indexForId := make(map[string]int)
 	for i, fileId := range fileIds {
 		strippedFileIds[i] = strings.TrimPrefix(fileId, "JDP:")
-		indexForId[strippedFileIds[i]] = i
 	}
 
-	// NOTE: the JDP search/by_file_ids/ endpoint (unofficial, undocumented!) only seems to
-	// NOTE: accept around 50 file IDs at a time, so we have to batch our requests
-
-	batchSize := 50
 	numBatches := len(strippedFileIds) / batchSize
 	if numBatches*batchSize < len(strippedFileIds) {
 		numBatches++
@@ -216,33 +209,60 @@ func (db *Database) Descriptors(orcid string, fileIds []string) ([]map[string]an
 			return nil, err
 		}
 
-		// get a de-duped list of descriptors
+		// get a de-duped list of descriptors (with JDP: prefixes reinstated on IDs)
 		batchDescriptors, err := descriptorsFromResponseBody(body, nil)
 		if err != nil {
 			return nil, err
 		}
 		descriptors = append(descriptors, batchDescriptors...)
 	}
+	return descriptors, nil
+}
+
+func (db *Database) Descriptors(orcid string, fileIds []string) ([]map[string]any, error) {
+	// NOTE: The JDP search/by_file_ids/ endpoint (unofficial, undocumented!) only seems to
+	// NOTE: accept around 50 file IDs at a time, so we have to batch our requests.
+	descriptors, err := db.fetchDescriptors(orcid, fileIds, 50)
+	if err != nil {
+		return nil, err
+	}
 
-	// reorder the descriptors to match that of the requested file IDs, and track file IDs that aren't
-	// matched to descriptors
 	descriptorsByFileId := make(map[string]map[string]any)
 	for _, descriptor := range descriptors {
 		descriptorsByFileId[descriptor["id"].(string)] = descriptor
 	}
 
-	// if any file IDs don't have corresponding descriptors, find out which ones and issue an error
+	// NOTE: Evidently sometimes the search/by_file_ids endpoint doesn't return all of the
+	// NOTE: relevant information (???), so we make a list of missing descriptors and attempt
+	// NOTE: to fetch them again. If that attempt fails, we emit an error.
 	if len(descriptorsByFileId) < len(fileIds) {
-		missingResources := make([]string, 0)
+		missingFileIds := make([]string, 0)
 		for _, fileId := range fileIds {
 			if _, found := descriptorsByFileId[fileId]; !found {
-				missingResources = append(missingResources, fileId)
+				missingFileIds = append(missingFileIds, fileId)
+			}
+		}
+		if recoveredDescriptors, err := db.fetchDescriptors(orcid, missingFileIds, 50); err == nil {
+			for _, descriptor := range recoveredDescriptors {
+				descriptorsByFileId[descriptor["id"].(string)] = descriptor
+			}
+			if len(recoveredDescriptors) < len(missingFileIds) { // didn't get them all!
+				missingFileIds = make([]string, 0)
+				for _, fileId := range fileIds {
+					if _, found := descriptorsByFileId[fileId]; !found {
+						missingFileIds = append(missingFileIds, fileId)
+					}
+				}
+			} else {
+				// got 'em!
+				missingFileIds = nil
 			}
 		}
-		if len(missingResources) > 0 {
+		if len(missingFileIds) > 0 {
+			slices.Sort(missingFileIds)
 			return nil, &databases.ResourcesNotFoundError{
 				Database:    "JDP",
-				ResourceIds: missingResources,
+				ResourceIds: missingFileIds,
 			}
 		}
 	}
@@ -258,9 +278,7 @@ func (db *Database) EndpointNames() []string {
 	return []string{db.EndpointName}
 }
 
-func (db *Database) StageFiles(orcid string, fileIds []string) (uuid.UUID, error) {
-	var xferId uuid.UUID
-
+func (db *Database) requestArchivedFiles(orcid string, fileIds []string, batchSize int) ([]int, error) {
 	// construct a POST request to restore archived files with the given IDs
 	type RestoreRequest struct {
 		Ids                []string `json:"ids"`
@@ -277,72 +295,116 @@ func (db *Database) StageFiles(orcid string, fileIds []string) (uuid.UUID, error
 		}
 	}
 
-	data, err := json.Marshal(RestoreRequest{
-		Ids:                fileIdsWithoutPrefix,
-		SendEmail:          false,
-		ApiVersion:         "2",
-		IncludePrivateData: 1, // we need this just in case!
-	})
-	if err != nil {
-		return xferId, err
+	numBatches := len(fileIds) / batchSize
+	if numBatches*batchSize < len(fileIds) {
+		numBatches += 1
 	}
 
-	// NOTE: The slash in the resource is all-important for POST requests to
-	// NOTE: the JDP!!
-	body, err := db.post("request_archived_files/", orcid, bytes.NewReader(data))
-	if err != nil {
-		switch e := err.(type) {
-		case *databases.ResourcesNotFoundError:
-			e.ResourceIds = fileIds
+	requestIds := make([]int, numBatches)
+	for i := range numBatches {
+		begin := i * batchSize
+		end := min((i+1)*batchSize, len(fileIdsWithoutPrefix))
+		data, err := json.Marshal(RestoreRequest{
+			Ids:                fileIdsWithoutPrefix[begin:end],
+			SendEmail:          false,
+			ApiVersion:         "2",
+			IncludePrivateData: 1, // we need this just in case!
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		// NOTE: The slash in the resource is all-important for POST requests to the JDP!!
+		body, err := db.post("request_archived_files/", orcid, bytes.NewReader(data))
+		if err != nil {
+			switch e := err.(type) {
+			case *databases.ResourcesNotFoundError:
+				e.ResourceIds = fileIds
+			}
+			return nil, err
+		}
+
+		type RestoreResponse struct {
+			RequestId int `json:"request_id"`
 		}
-		return xferId, err
-	}
 
-	type RestoreResponse struct {
-		RequestId int `json:"request_id"`
+		var jdpResp RestoreResponse
+		err = json.Unmarshal(body, &jdpResp)
+		if err != nil {
+			return nil, err
+		}
+		requestIds[i] = jdpResp.RequestId
 	}
 
-	var jdpResp RestoreResponse
-	err = json.Unmarshal(body, &jdpResp)
+	return requestIds, nil
+}
+
+func (db *Database) StageFiles(orcid string, fileIds []string) (uuid.UUID, error) {
+	var xferId uuid.UUID
+
+	// NOTE: the relevant endpoint seems to return a 404 whenever it gets too many file IDs,
+	// NOTE: so we batch requests in sets of 1000
+	requestIds, err := db.requestArchivedFiles(orcid, fileIds, 1000)
 	if err != nil {
 		return xferId, err
 	}
-	slog.Debug(fmt.Sprintf("Requested %d archived files from JDP (request ID: %d)",
-		len(fileIds), jdpResp.RequestId))
+
+	slog.Debug(fmt.Sprintf("Requested %d archived files from JDP (request IDs: %v)",
+		len(fileIds), requestIds))
 	xferId = uuid.New()
 	db.StagingRequests[xferId] = StagingRequest{
-		Id:   jdpResp.RequestId,
+		Ids:  requestIds,
 		Time: time.Now(),
 	}
 	return xferId, err
 }
 
 func (db *Database) StagingStatus(id uuid.UUID) (databases.StagingStatus, error) {
+	statusForString := map[string]databases.StagingStatus{
+		"new":     databases.StagingStatusActive,
+		"pending": databases.StagingStatusActive,
+		"ready":   databases.StagingStatusSucceeded,
+		"failed":  databases.StagingStatusFailed,
+	}
 	db.pruneStagingRequests()
 	if request, found := db.StagingRequests[id]; found {
-		resource := fmt.Sprintf("request_archived_files/requests/%d", request.Id)
-		body, err := db.get(resource, url.Values{})
-		if err != nil {
-			return databases.StagingStatusUnknown, err
-		}
-		type JDPResult struct {
-			Status string `json:"status"` // "new", "pending", or "ready"
-		}
-		var jdpResult JDPResult
-		err = json.Unmarshal(body, &jdpResult)
-		if err != nil {
-			return databases.StagingStatusUnknown, err
-		}
-		statusForString := map[string]databases.StagingStatus{
-			"new":     databases.StagingStatusActive,
-			"pending": databases.StagingStatusActive,
-			"ready":   databases.StagingStatusSucceeded,
-		}
-		if status, ok := statusForString[jdpResult.Status]; ok {
-			slog.Debug(fmt.Sprintf("Queried JDP for staging status of transfer with staging ID %s (request ID: %d): %s", id, request.Id, jdpResult.Status))
-			return status, nil
+		var status databases.StagingStatus
+		var statusStr string
+		for _, requestId := range request.Ids {
+			resource := fmt.Sprintf("request_archived_files/requests/%d", requestId)
+			body, err := db.get(resource, url.Values{})
+			if err != nil {
+				return databases.StagingStatusUnknown, err
+			}
+			type JDPResult struct {
+				Status string `json:"status"` // "new", "pending", "ready", or "failed"
+			}
+			var jdpResult JDPResult
+			err = json.Unmarshal(body, &jdpResult)
+			if err != nil {
+				return databases.StagingStatusUnknown, err
+			}
+			if requestStatus, ok := statusForString[jdpResult.Status]; ok {
+				if status == databases.StagingStatusUnknown { // first status encountered
+					status = requestStatus
+					statusStr = jdpResult.Status
+				} else {
+					if requestStatus != status { // status update
+						if requestStatus != databases.StagingStatusSucceeded {
+							status = requestStatus
+							statusStr = jdpResult.Status
+						}
+					}
+				}
+				if status == databases.StagingStatusFailed { // one failure sinks them all
+					break
+				}
+			} else {
+				return databases.StagingStatusUnknown, fmt.Errorf("unrecognized JDP staging status string: %s", jdpResult.Status)
+			}
 		}
-		return databases.StagingStatusUnknown, fmt.Errorf("unrecognized JDP staging status string: %s", jdpResult.Status)
+		slog.Debug(fmt.Sprintf("Queried JDP for staging status of transfer with staging ID %s (request IDs: %v): %s", id, request.Ids, statusStr))
+		return status, nil
 	} else {
 		slog.Info(fmt.Sprintf("No staging request found for transfer with staging ID %s", id.String()))
 		return databases.StagingStatusUnknown, nil
@@ -660,10 +722,6 @@ func (db *Database) post(resource, orcid string, body io.Reader) ([]byte, error)
 	case 200, 201, 204:
 		defer resp.Body.Close()
 		return io.ReadAll(resp.Body)
-	case 404:
-		return nil, &databases.ResourcesNotFoundError{
-			Database: "JDP",
-		}
 	case 503:
 		return nil, &databases.UnavailableError{
 			Database: "jdp",
@@ -855,7 +913,7 @@ func (db *Database) pruneStagingRequests() {
 	for uuid, request := range db.StagingRequests {
 		requestAge := time.Since(request.Time)
 		if requestAge > db.DeleteAfter {
-			slog.Debug(fmt.Sprintf("Pruning staging request with staging ID %s (request ID: %d): age (%s) exceeds limit (%s)", uuid.String(), request.Id, requestAge.String(), db.DeleteAfter.String()))
+			slog.Debug(fmt.Sprintf("Pruning staging request with staging ID %s (request IDs: %v): age (%s) exceeds limit (%s)", uuid.String(), request.Ids, requestAge.String(), db.DeleteAfter.String()))
 			delete(db.StagingRequests, uuid)
 		}
 	}

etc/irods/pyproject.toml

@@ -6,4 +6,5 @@ requires-python = ">=3.10"
 dependencies = [
     "boto3>=1.41.5",
     "python-irodsclient>=3.2.0",
+    "urllib3>=2.7.0",
 ]

etc/irods/uv.lock

@@ -6,11 +6,11 @@ toolchain go1.24.0
 
 require (
 	github.com/StalkR/hsts v1.0.3
-	github.com/aws/aws-sdk-go-v2 v1.39.6
+	github.com/aws/aws-sdk-go-v2 v1.41.5
 	github.com/aws/aws-sdk-go-v2/config v1.31.17
 	github.com/aws/aws-sdk-go-v2/credentials v1.18.21
 	github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.20.4
-	github.com/aws/aws-sdk-go-v2/service/s3 v1.90.0
+	github.com/aws/aws-sdk-go-v2/service/s3 v1.97.3
 	github.com/danielgtaylor/huma/v2 v2.27.0
 	github.com/fernet/fernet-go v0.0.0-20240119011108-303da6aec611
 	github.com/frictionlessdata/datapackage-go v1.0.4
@@ -26,14 +26,14 @@ require (
 require (
 	github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.8 // indirect
 	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.13 // indirect
-	github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.13 // indirect
-	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.13 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.21 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.7.21 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/ini v1.8.4 // indirect
-	github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.13 // indirect
-	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.3 // indirect
-	github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.4 // indirect
-	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.13 // indirect
-	github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.13 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/v4a v1.4.22 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.13.7 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.9.13 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.13.21 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.19.21 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sso v1.30.1 // indirect
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.35.5 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sts v1.39.1 // indirect