001/*-
002 * #%L
003 * HAPI FHIR JPA Server - Batch2 Task Processor
004 * %%
005 * Copyright (C) 2014 - 2024 Smile CDR, Inc.
006 * %%
007 * Licensed under the Apache License, Version 2.0 (the "License");
008 * you may not use this file except in compliance with the License.
009 * You may obtain a copy of the License at
010 *
011 *      http://www.apache.org/licenses/LICENSE-2.0
012 *
013 * Unless required by applicable law or agreed to in writing, software
014 * distributed under the License is distributed on an "AS IS" BASIS,
015 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
016 * See the License for the specific language governing permissions and
017 * limitations under the License.
018 * #L%
019 */
020package ca.uhn.fhir.batch2.coordinator;
021
022import ca.uhn.fhir.batch2.api.IJobPartitionProvider;
023import ca.uhn.fhir.batch2.jobs.parameters.PartitionedUrl;
024import ca.uhn.fhir.context.FhirContext;
025import ca.uhn.fhir.interceptor.model.RequestPartitionId;
026import ca.uhn.fhir.jpa.partition.IRequestPartitionHelperSvc;
027import ca.uhn.fhir.jpa.searchparam.MatchUrlService;
028import ca.uhn.fhir.jpa.searchparam.ResourceSearch;
029import ca.uhn.fhir.rest.api.server.RequestDetails;
030
031import java.util.ArrayList;
032import java.util.LinkedHashSet;
033import java.util.List;
034import java.util.Set;
035import java.util.stream.Collectors;
036
037/**
038 * Default implementation which provides the {@link PartitionedUrl} list for a certain operation request.
039 */
040public class DefaultJobPartitionProvider implements IJobPartitionProvider {
041        protected final IRequestPartitionHelperSvc myRequestPartitionHelper;
042        protected FhirContext myFhirContext;
043        private MatchUrlService myMatchUrlService;
044
045        public DefaultJobPartitionProvider(IRequestPartitionHelperSvc theRequestPartitionHelperSvc) {
046                myRequestPartitionHelper = theRequestPartitionHelperSvc;
047        }
048
049        public DefaultJobPartitionProvider(
050                        FhirContext theFhirContext,
051                        IRequestPartitionHelperSvc theRequestPartitionHelperSvc,
052                        MatchUrlService theMatchUrlService) {
053                myFhirContext = theFhirContext;
054                myRequestPartitionHelper = theRequestPartitionHelperSvc;
055                myMatchUrlService = theMatchUrlService;
056        }
057
058        public List<RequestPartitionId> getPartitions(RequestDetails theRequestDetails, String theOperation) {
059                RequestPartitionId partitionId = myRequestPartitionHelper.determineReadPartitionForRequestForServerOperation(
060                                theRequestDetails, theOperation);
061                return List.of(partitionId);
062        }
063
064        @Override
065        public List<PartitionedUrl> getPartitionedUrls(RequestDetails theRequestDetails, List<String> theUrls) {
066                List<String> urls = theUrls;
067
068                // if the url list is empty, use all the supported resource types to build the url list
069                // we can go back to no url scenario if all resource types point to the same partition
070                if (theUrls == null || theUrls.isEmpty()) {
071                        urls = myFhirContext.getResourceTypes().stream()
072                                        .map(resourceType -> resourceType + "?")
073                                        .collect(Collectors.toList());
074                }
075
076                // determine the partition associated with each of the urls
077                List<PartitionedUrl> partitionedUrls = new ArrayList<>();
078                for (String s : urls) {
079                        ResourceSearch resourceSearch = myMatchUrlService.getResourceSearch(s);
080                        RequestPartitionId partitionId = myRequestPartitionHelper.determineReadPartitionForRequestForSearchType(
081                                        theRequestDetails, resourceSearch.getResourceName(), resourceSearch.getSearchParameterMap());
082                        partitionedUrls.add(new PartitionedUrl().setUrl(s).setRequestPartitionId(partitionId));
083                }
084
085                // handle (bulk) system operations that are typically configured with RequestPartitionId.allPartitions()
086                // populate the actual list of all partitions, if that is supported
087                Set<RequestPartitionId> allPartitions = new LinkedHashSet<>(getAllPartitions());
088                List<PartitionedUrl> retVal = new ArrayList<>();
089                for (PartitionedUrl partitionedUrl : partitionedUrls) {
090                        String url = partitionedUrl.getUrl();
091                        RequestPartitionId partitionId = partitionedUrl.getRequestPartitionId();
092                        if (partitionId != null && partitionId.isAllPartitions() && !allPartitions.isEmpty()) {
093                                allPartitions.stream()
094                                                .map(p -> (new PartitionedUrl().setUrl(url).setRequestPartitionId(p)))
095                                                .forEach(retVal::add);
096                        } else {
097                                retVal.add(partitionedUrl);
098                        }
099                }
100
101                return retVal;
102        }
103
104        public List<RequestPartitionId> getAllPartitions() {
105                return List.of(RequestPartitionId.allPartitions());
106        }
107}