001/*-
002 * #%L
003 * HAPI FHIR JPA Server
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.jpa.search;
021
022import ca.uhn.fhir.interceptor.model.RequestPartitionId;
023import ca.uhn.fhir.jpa.dao.ISearchBuilder;
024import ca.uhn.fhir.jpa.entity.SearchTypeEnum;
025import ca.uhn.fhir.jpa.model.dao.JpaPid;
026import ca.uhn.fhir.jpa.model.search.SearchStatusEnum;
027import ca.uhn.fhir.jpa.search.builder.tasks.SearchTask;
028import ca.uhn.fhir.jpa.util.QueryParameterUtils;
029import ca.uhn.fhir.model.api.ResourceMetadataKeyEnum;
030import ca.uhn.fhir.model.valueset.BundleEntrySearchModeEnum;
031import ca.uhn.fhir.rest.api.server.RequestDetails;
032import ca.uhn.fhir.rest.server.method.ResponsePage;
033import jakarta.annotation.Nonnull;
034import org.hl7.fhir.instance.model.api.IBaseResource;
035import org.slf4j.Logger;
036import org.slf4j.LoggerFactory;
037
038import java.util.List;
039import java.util.Set;
040import java.util.stream.Collectors;
041
042public class PersistedJpaSearchFirstPageBundleProvider extends PersistedJpaBundleProvider {
043        private static final Logger ourLog = LoggerFactory.getLogger(PersistedJpaSearchFirstPageBundleProvider.class);
044        private final SearchTask mySearchTask;
045
046        @SuppressWarnings("rawtypes")
047        private final ISearchBuilder mySearchBuilder;
048
049        /**
050         * Constructor
051         */
052        @SuppressWarnings("rawtypes")
053        public PersistedJpaSearchFirstPageBundleProvider(
054                        SearchTask theSearchTask,
055                        ISearchBuilder theSearchBuilder,
056                        RequestDetails theRequest,
057                        RequestPartitionId theRequestPartitionId) {
058                super(theRequest, theSearchTask.getSearch());
059
060                assert getSearchEntity().getSearchType() != SearchTypeEnum.HISTORY;
061
062                mySearchTask = theSearchTask;
063                mySearchBuilder = theSearchBuilder;
064                super.setRequestPartitionId(theRequestPartitionId);
065        }
066
067        @Nonnull
068        @Override
069        public List<IBaseResource> getResources(
070                        int theFromIndex, int theToIndex, @Nonnull ResponsePage.ResponsePageBuilder thePageBuilder) {
071                ensureSearchEntityLoaded();
072                QueryParameterUtils.verifySearchHasntFailedOrThrowInternalErrorException(getSearchEntity());
073
074                mySearchTask.awaitInitialSync();
075
076                // request 1 more than we need to, in order to know if there are extra values
077                ourLog.trace("Fetching search resource PIDs from task: {}", mySearchTask.getClass());
078                final List<JpaPid> pids = mySearchTask.getResourcePids(theFromIndex, theToIndex + 1);
079                ourLog.trace("Done fetching search resource PIDs");
080
081                int countOfPids = pids.size();
082
083                int maxSize = Math.min(theToIndex - theFromIndex, countOfPids);
084                thePageBuilder.setTotalRequestedResourcesFetched(countOfPids);
085
086                RequestPartitionId requestPartitionId = getRequestPartitionId();
087
088                List<JpaPid> firstBatch = pids.subList(0, maxSize);
089                List<IBaseResource> retVal = myTxService
090                                .withRequest(myRequest)
091                                .withRequestPartitionId(requestPartitionId)
092                                .execute(() -> toResourceList(mySearchBuilder, firstBatch, thePageBuilder));
093
094                long totalCountWanted = theToIndex - theFromIndex;
095                long totalCountMatch = (int) retVal.stream().filter(t -> !isInclude(t)).count();
096
097                if (totalCountMatch < totalCountWanted) {
098                        if (getSearchEntity().getStatus() == SearchStatusEnum.PASSCMPLET
099                                        || ((getSearchEntity().getStatus() == SearchStatusEnum.FINISHED
100                                                        && getSearchEntity().getNumFound() >= theToIndex))) {
101
102                                /*
103                                 * This is a bit of complexity to account for the possibility that
104                                 * the consent service has filtered some results.
105                                 */
106                                Set<String> existingIds = retVal.stream()
107                                                .map(t -> t.getIdElement().getValue())
108                                                .filter(t -> t != null)
109                                                .collect(Collectors.toSet());
110
111                                long remainingWanted = totalCountWanted - totalCountMatch;
112                                long fromIndex = theToIndex - remainingWanted;
113                                ResponsePage.ResponsePageBuilder pageBuilder = new ResponsePage.ResponsePageBuilder();
114                                pageBuilder.setBundleProvider(this);
115                                List<IBaseResource> remaining = super.getResources((int) fromIndex, theToIndex, pageBuilder);
116                                remaining.forEach(t -> {
117                                        if (!existingIds.contains(t.getIdElement().getValue())) {
118                                                retVal.add(t);
119                                        }
120                                });
121                                thePageBuilder.combineWith(pageBuilder);
122                        }
123                }
124                ourLog.trace("Loaded resources to return");
125
126                return retVal;
127        }
128
129        private boolean isInclude(IBaseResource theResource) {
130                BundleEntrySearchModeEnum searchMode = ResourceMetadataKeyEnum.ENTRY_SEARCH_MODE.get(theResource);
131                return BundleEntrySearchModeEnum.INCLUDE.equals(searchMode);
132        }
133
134        @Override
135        public Integer size() {
136                ourLog.trace("size() - Waiting for initial sync");
137                Integer size = mySearchTask.awaitInitialSync();
138                ourLog.trace("size() - Finished waiting for local sync");
139
140                ensureSearchEntityLoaded();
141                QueryParameterUtils.verifySearchHasntFailedOrThrowInternalErrorException(getSearchEntity());
142                if (size != null) {
143                        return size;
144                }
145                return super.size();
146        }
147}