001/*-
002 * #%L
003 * HAPI FHIR JPA Server
004 * %%
005 * Copyright (C) 2014 - 2023 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.Search;
025import ca.uhn.fhir.jpa.entity.SearchTypeEnum;
026import ca.uhn.fhir.jpa.model.dao.JpaPid;
027import ca.uhn.fhir.jpa.model.search.SearchStatusEnum;
028import ca.uhn.fhir.jpa.search.builder.tasks.SearchTask;
029import ca.uhn.fhir.jpa.util.QueryParameterUtils;
030import ca.uhn.fhir.model.api.ResourceMetadataKeyEnum;
031import ca.uhn.fhir.model.valueset.BundleEntrySearchModeEnum;
032import ca.uhn.fhir.rest.api.server.RequestDetails;
033import org.hl7.fhir.instance.model.api.IBaseResource;
034import org.slf4j.Logger;
035import org.slf4j.LoggerFactory;
036
037import javax.annotation.Nonnull;
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        private final ISearchBuilder mySearchBuilder;
046
047        /**
048         * Constructor
049         */
050        public PersistedJpaSearchFirstPageBundleProvider(Search theSearch, SearchTask theSearchTask, ISearchBuilder theSearchBuilder, RequestDetails theRequest, RequestPartitionId theRequestPartitionId) {
051                super(theRequest, theSearch.getUuid());
052
053                assert theSearch.getSearchType() != SearchTypeEnum.HISTORY;
054
055                setSearchEntity(theSearch);
056                mySearchTask = theSearchTask;
057                mySearchBuilder = theSearchBuilder;
058                super.setRequestPartitionId(theRequestPartitionId);
059        }
060
061        @Nonnull
062        @Override
063        public List<IBaseResource> getResources(int theFromIndex, int theToIndex) {
064                ensureSearchEntityLoaded();
065                QueryParameterUtils.verifySearchHasntFailedOrThrowInternalErrorException(getSearchEntity());
066
067                mySearchTask.awaitInitialSync();
068
069                ourLog.trace("Fetching search resource PIDs from task: {}", mySearchTask.getClass());
070                final List<JpaPid> pids = mySearchTask.getResourcePids(theFromIndex, theToIndex);
071                ourLog.trace("Done fetching search resource PIDs");
072
073                RequestPartitionId requestPartitionId = getRequestPartitionId();
074
075                List<IBaseResource> retVal = myTxService
076                        .withRequest(myRequest)
077                        .withRequestPartitionId(requestPartitionId)
078                        .execute(() -> toResourceList(mySearchBuilder, pids));
079
080                long totalCountWanted = theToIndex - theFromIndex;
081                long totalCountMatch = (int) retVal
082                        .stream()
083                        .filter(t -> !isInclude(t))
084                        .count();
085
086                if (totalCountMatch < totalCountWanted) {
087                        if (getSearchEntity().getStatus() == SearchStatusEnum.PASSCMPLET
088                                || ((getSearchEntity().getStatus() == SearchStatusEnum.FINISHED && getSearchEntity().getNumFound() >= theToIndex))) {
089
090                                /*
091                                 * This is a bit of complexity to account for the possibility that
092                                 * the consent service has filtered some results.
093                                 */
094                                Set<String> existingIds = retVal
095                                        .stream()
096                                        .map(t -> t.getIdElement().getValue())
097                                        .filter(t -> t != null)
098                                        .collect(Collectors.toSet());
099
100                                long remainingWanted = totalCountWanted - totalCountMatch;
101                                long fromIndex = theToIndex - remainingWanted;
102                                List<IBaseResource> remaining = super.getResources((int) fromIndex, theToIndex);
103                                remaining.forEach(t -> {
104                                        if (!existingIds.contains(t.getIdElement().getValue())) {
105                                                retVal.add(t);
106                                        }
107                                });
108                        }
109                }
110                ourLog.trace("Loaded resources to return");
111
112                return retVal;
113        }
114
115        private boolean isInclude(IBaseResource theResource) {
116                BundleEntrySearchModeEnum searchMode = ResourceMetadataKeyEnum.ENTRY_SEARCH_MODE.get(theResource);
117                return BundleEntrySearchModeEnum.INCLUDE.equals(searchMode);
118        }
119
120        @Override
121        public Integer size() {
122                ourLog.trace("size() - Waiting for initial sync");
123                Integer size = mySearchTask.awaitInitialSync();
124                ourLog.trace("size() - Finished waiting for local sync");
125
126                ensureSearchEntityLoaded();
127                QueryParameterUtils.verifySearchHasntFailedOrThrowInternalErrorException(getSearchEntity());
128                if (size != null) {
129                        return size;
130                }
131                return super.size();
132        }
133
134}