Source code

001/*
002    Licensed to the Apache Software Foundation (ASF) under one
003    or more contributor license agreements.  See the NOTICE file
004    distributed with this work for additional information
005    regarding copyright ownership.  The ASF licenses this file
006    to you under the Apache License, Version 2.0 (the
007    "License"); you may not use this file except in compliance
008    with the License.  You may obtain a copy of the License at
009
010       http://www.apache.org/licenses/LICENSE-2.0
011
012    Unless required by applicable law or agreed to in writing,
013    software distributed under the License is distributed on an
014    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
015    KIND, either express or implied.  See the License for the
016    specific language governing permissions and limitations
017    under the License.    
018 */
019package org.apache.wiki.search.kendra;
020
021import com.amazonaws.services.kendra.AWSkendra;
022import com.amazonaws.services.kendra.AWSkendraClientBuilder;
023import com.amazonaws.services.kendra.model.*;
024import com.amazonaws.util.IOUtils;
025import com.google.gson.Gson;
026import com.google.gson.GsonBuilder;
027import com.google.gson.reflect.TypeToken;
028import org.apache.commons.io.FilenameUtils;
029import org.apache.commons.lang3.Strings;
030import org.apache.logging.log4j.LogManager;
031import org.apache.logging.log4j.Logger;
032import org.apache.wiki.InternalWikiException;
033import org.apache.wiki.WatchDog;
034import org.apache.wiki.WikiBackgroundThread;
035import org.apache.wiki.api.core.Attachment;
036import org.apache.wiki.api.core.Context;
037import org.apache.wiki.api.core.Engine;
038import org.apache.wiki.api.core.Page;
039import org.apache.wiki.api.exceptions.NoRequiredPropertyException;
040import org.apache.wiki.api.exceptions.ProviderException;
041import org.apache.wiki.api.providers.PageProvider;
042import org.apache.wiki.api.search.SearchResult;
043import org.apache.wiki.api.spi.Wiki;
044import org.apache.wiki.attachment.AttachmentManager;
045import org.apache.wiki.auth.AuthorizationManager;
046import org.apache.wiki.auth.permissions.PagePermission;
047import org.apache.wiki.pages.PageManager;
048import org.apache.wiki.search.SearchProvider;
049import org.apache.wiki.util.TextUtil;
050
051import java.io.IOException;
052import java.io.InputStream;
053import java.io.InputStreamReader;
054import java.lang.reflect.Type;
055import java.nio.ByteBuffer;
056import java.nio.charset.StandardCharsets;
057import java.util.ArrayList;
058import java.util.Collection;
059import java.util.Collections;
060import java.util.HashMap;
061import java.util.List;
062import java.util.Map;
063import java.util.Properties;
064
065import static java.lang.String.format;
066
067/**
068 * Search provider that implements {link SearchProvider} using AWS Kendra for
069 * indexing. Note that we are using a Custom DataSource which limits the
070 * attributes that can be uploaded / searched for each page (as per
071 * https://docs.aws.amazon.com/kendra/latest/dg/custom-attributes.html). This
072 * could be overcome by using an S3 bucket where any custom attributes can be
073 * added.
074 *
075 * @since 2.11.0
076 */
077public class KendraSearchProvider implements SearchProvider {
078
079    private static final Logger LOG = LogManager.getLogger( KendraSearchProvider.class );
080    private Engine engine;
081    private Properties properties;
082    private Map< String, Object > contentTypes;
083    private AWSkendra kendra;
084    private String indexName;
085    private String indexId;
086    private String dataSourceName;
087    private String dataSourceId;
088
089    private final List< Page > updates = Collections.synchronizedList( new ArrayList<>() );
090
091    private static final String PROP_KENDRA_INDEX_NAME = "jspwiki.kendra.indexName";
092    private static final String PROP_KENDRA_DATA_SOURCE_NAME = "jspwiki.kendra.dataSourceName";
093    private static final String PROP_KENDRA_INDEXDELAY = "jspwiki.kendra.indexdelay";
094    private static final String PROP_KENDRA_INITIALDELAY = "jspwiki.kendra.initialdelay";
095
096    public KendraSearchProvider() {
097    }
098
099    /**
100     * {@inheritDoc}
101     */
102    @Override
103    public void initialize( final Engine engine, final Properties properties ) throws NoRequiredPropertyException, IOException {
104        this.engine = engine;
105        this.properties = properties;
106        this.contentTypes = getContentTypes();
107
108        setKendra( buildClient() );
109
110        this.indexName = TextUtil.getRequiredProperty( this.properties, PROP_KENDRA_INDEX_NAME );
111        this.dataSourceName = TextUtil.getRequiredProperty( this.properties, PROP_KENDRA_DATA_SOURCE_NAME );
112        final int initialDelay = TextUtil.getIntegerProperty( this.properties, PROP_KENDRA_INITIALDELAY, KendraUpdater.INITIAL_DELAY );
113        final int indexDelay = TextUtil.getIntegerProperty( this.properties, PROP_KENDRA_INDEXDELAY, KendraUpdater.INDEX_DELAY );
114
115        // Start the Kendra update thread, which waits first for a little while
116        // before starting to go through the "pages that need updating".
117        if ( initialDelay >= 0 ) {
118            final KendraUpdater updater = new KendraUpdater( engine, this, initialDelay, indexDelay );
119            updater.start();
120        }
121    }
122
123    private Map< String, Object > getContentTypes() {
124        final Gson gson = new GsonBuilder().create();
125        try ( final InputStream in = KendraSearchProvider.class.getResourceAsStream( "content_types.json" ) ) {
126            if ( in != null ) {
127                final Type collectionType = new TypeToken< HashMap< String, Object > >() {
128                }.getType();
129                return gson.fromJson( new InputStreamReader( in ), collectionType );
130            }
131        } catch ( final IOException e ) {
132            LOG.error( "Unable to load default propertyfile 'content_types.json': {}", e.getMessage(), e );
133        }
134        return null;
135    }
136
137    /**
138     * {@inheritDoc}
139     */
140    @Override
141    public String getProviderInfo() {
142        return "KendraSearchProvider";
143    }
144
145    /**
146     * {@inheritDoc}
147     */
148    @Override
149    public void pageRemoved( final Page page ) {
150        final String pageName = page.getName();
151        final BatchDeleteDocumentRequest request = new BatchDeleteDocumentRequest().withIndexId( indexId )
152                .withDocumentIdList( pageName );
153        final BatchDeleteDocumentResult result = getKendra().batchDeleteDocument( request );
154        if (result.getFailedDocuments().isEmpty()) {
155            LOG.debug( "Page '{}' was removed from index", pageName );
156        } else {
157            LOG.error( "Failed to remove Page '{}' from index", pageName );
158        }
159    }
160
161    /**
162     * {@inheritDoc}
163     */
164    @Override
165    public void reindexPage( final Page page ) {
166        if ( page != null ) {
167            updates.add( page );
168            LOG.debug( format( "Scheduling page '%s' for indexing ...", page.getName() ) );
169        }
170    }
171
172    /**
173     * {@inheritDoc}
174     */
175    @Override
176    public Collection< SearchResult > findPages( final String query, final Context wikiContext ) throws ProviderException, IOException {
177        final QueryRequest request = new QueryRequest().withIndexId( indexId ).withQueryText( query );
178        final List< QueryResultItem > items;
179        try {
180            items = getKendra().query( request ).getResultItems();
181        } catch ( final ThrottlingException e ) {
182            LOG.error( "ThrottlingException. Skipping..." );
183            return new ArrayList<>();
184        }
185        final List< SearchResult > searchResults = new ArrayList<>( items.size() );
186        final AuthorizationManager mgr = engine.getManager( AuthorizationManager.class );
187
188        for ( final QueryResultItem item : items ) {
189            switch( QueryResultType.fromValue( item.getType() ) ) {
190                case DOCUMENT:
191                    final String documentId = item.getDocumentId();
192                    final String documentExcerpt = item.getDocumentExcerpt().getText();
193                    final String scoreConfidence = item.getScoreAttributes().getScoreConfidence();
194                    final Page page = this.engine.getManager( PageManager.class ).getPage( documentId, PageProvider.LATEST_VERSION );
195                    if ( page != null ) {
196                        final PagePermission pp = new PagePermission( page, PagePermission.VIEW_ACTION );
197                        if ( mgr.checkPermission( wikiContext.getWikiSession(), pp ) ) {
198                            final SearchResult searchResult = new SearchResultImpl( page, confidence2score( scoreConfidence ),
199                                    new String[]{ documentExcerpt } );
200                            searchResults.add( searchResult );
201                        } else {
202                            LOG.error( format( "Page '%s' is not accessible", documentId ) );
203                        }
204                    } else {
205                        LOG.error(
206                                format( "Kendra found a result page '%s' that could not be loaded, removing from index", documentId ) );
207                        pageRemoved( Wiki.contents().page( this.engine, documentId ) );
208                    }
209                    break;
210                default:
211                    LOG.error( format( "Unknown query result type: %s", item.getType() ) );
212            }
213        }
214        return searchResults;
215    }
216
217    /**
218     * This method initialize the AWS Kendra Index and Datasources to be used.
219     */
220    public void initializeIndexAndDataSource() {
221        this.indexId = getIndexId( indexName );
222        if ( this.indexId == null ) {
223            final String message = format( "Index '%s' does not exist", indexName );
224            LOG.error( message );
225            throw new IllegalArgumentException( message );
226        }
227        this.dataSourceId = getDatasourceId( this.indexId, dataSourceName );
228        if ( this.dataSourceId == null ) {
229            final String message = format( "Datasource '%s' does not exist in index %s", dataSourceName, indexName );
230            LOG.error( message );
231            throw new IllegalArgumentException( message );
232        }
233    }
234
235    /**
236     * Given an Kendra's Index name, returns the corresponding Index Id, or
237     * {@code null} if it does not exists
238     *
239     * @param indexName the name of the index to look up
240     * @return the index id or {@code null}
241     */
242    private String getIndexId( final String indexName ) {
243        ListIndicesRequest request = new ListIndicesRequest();
244        ListIndicesResult result = getKendra().listIndices( request );
245        String nextToken = "";
246        while ( nextToken != null ) {
247            final List< IndexConfigurationSummary > items = result.getIndexConfigurationSummaryItems();
248            if ( items == null || items.isEmpty() ) {
249                return null;
250            }
251            for ( final IndexConfigurationSummary item : items ) {
252                if ( Strings.CS.equals( item.getName(), indexName ) ) {
253                    return item.getId();
254                }
255            }
256            nextToken = result.getNextToken();
257            request = new ListIndicesRequest().withNextToken( result.getNextToken() );
258            result = getKendra().listIndices( request );
259        }
260        return null;
261    }
262
263    /**
264     * Given an Kendra's Datasource name, returns the corresponding Datasource Id,
265     * or {@code null} if it does not exists
266     *
267     * @param dataSourceName the name of the datasource to look up
268     * @return the datasource id or {@code null}
269     */
270    private String getDatasourceId( final String indexId, final String dataSourceName ) {
271        ListDataSourcesRequest request = new ListDataSourcesRequest().withIndexId( indexId );
272        ListDataSourcesResult result = getKendra().listDataSources( request );
273        String nextToken = "";
274        while ( nextToken != null ) {
275            final List< DataSourceSummary > items = result.getSummaryItems();
276            if( items == null || items.isEmpty() ) {
277                return null;
278            }
279
280            for( final DataSourceSummary item : items ) {
281                if( Strings.CS.equals( item.getName(), dataSourceName ) ) {
282                    return item.getId();
283                }
284            }
285            nextToken = result.getNextToken();
286            request = new ListDataSourcesRequest().withNextToken( result.getNextToken() );
287            result = getKendra().listDataSources( request );
288        }
289        return null;
290    }
291
292    /*
293     * Converts a SCORE Confidence from Kendra to an "equivalent" integer score
294     */
295    private int confidence2score( final String scoreConfidence ) {
296        return switch( ScoreConfidence.fromValue( scoreConfidence ) ) {
297            case VERY_HIGH -> 100;
298            case HIGH -> 75;
299            case MEDIUM -> 50;
300            case LOW -> 25;
301            default -> 0;
302        };
303    }
304
305    /**
306     * This method re-index all the pages found in the Wiki. It is mainly used at
307     * startup.
308     *
309     * @throws IOException in case some page can not be read
310     */
311    private void doFullReindex() throws IOException {
312        try {
313            final Collection< Page > pages = engine.getManager( PageManager.class ).getAllPages();
314            if ( pages.isEmpty() ) {
315                return;
316            }
317            LOG.debug( format( "Indexing all %d pages. Please wait ...", pages.size() ) );
318            final String executionId = startExecution();
319            for ( final Page page : pages ) {
320                // Since I do not want to handle the size limit
321                // (https://docs.aws.amazon.com/goto/WebAPI/kendra-2019-02-03/BatchPutDocument)
322                // uploading documents one at a time
323                indexOnePage( page, executionId );
324            }
325        } catch ( final ProviderException e ) {
326            LOG.error( e.getMessage() );
327            throw new IOException( e );
328        } finally {
329            stopExecution();
330        }
331    }
332
333    /**
334     * This method re-index all pages marked as updated. It is used to periodically
335     * index pages that have been modified
336     */
337    private void doPartialReindex() {
338        if ( updates.isEmpty() ) {
339            return;
340        }
341        LOG.debug( "Indexing updated pages. Please wait ..." );
342        final String executionId = startExecution();
343        synchronized ( updates ) {
344            try {
345                while (!updates.isEmpty()) {
346                    indexOnePage( updates.remove( 0 ), executionId );
347                }
348            } finally {
349                stopExecution();
350            }
351        }
352    }
353
354    /**
355     * Returns an ExecutiuonId that is required to keep track of the modifed
356     * documents
357     *
358     * @return The execution id
359     */
360    private String startExecution() {
361        final StartDataSourceSyncJobRequest request = new StartDataSourceSyncJobRequest().withIndexId( indexId )
362                .withId( dataSourceId );
363        final StartDataSourceSyncJobResult result = getKendra().startDataSourceSyncJob( request );
364        return result.getExecutionId();
365    }
366
367    /**
368     * Stop the execution for the given index Id and DataSource Id.
369     */
370    private void stopExecution() {
371        final StopDataSourceSyncJobRequest request = new StopDataSourceSyncJobRequest().withIndexId( indexId ).withId( dataSourceId );
372        getKendra().stopDataSourceSyncJob( request );
373    }
374
375    /**
376     * Index on single {@link Page} into the Kendra Index
377     *
378     * @param page        the {@link Page} to index
379     * @param executionId The Execution Id
380     */
381    private void indexOnePage( final Page page, final String executionId ) {
382        final String pageName = page.getName();
383        try {
384            final Document document = newDocument( page, executionId );
385            final BatchPutDocumentRequest request = new BatchPutDocumentRequest().withIndexId( indexId )
386                    .withDocuments( document );
387            final BatchPutDocumentResult result = getKendra().batchPutDocument( request );
388            if (result.getFailedDocuments().isEmpty()) {
389                LOG.info( format( "Successfully indexed Page '%s' as %s", page.getName(), document.getContentType() ) );
390            } else {
391                for ( final BatchPutDocumentResponseFailedDocument failedDocument : result.getFailedDocuments() ) {
392                    LOG.error( format( "Failed to index Page '%s': %s", failedDocument.getId(), failedDocument.getErrorMessage() ) );
393                }
394            }
395        } catch ( final IOException e ) {
396            LOG.error( format( "Failed to index Page '%s': %s", pageName, e.getMessage() ) );
397        }
398    }
399
400
401    /**
402     * Given a {@link Page}, returns the corresponding Kendra {@link Document}.
403     *
404     * @param page        the {@link Page} to be indexed
405     * @param executionId an execution id to identify when the {@link Page} was
406     *                    indexed for the last time.
407     * @return a {@link Document} containing the searchable attributes.
408     * @throws IOException if the {@link Page}'s {@link Attachment} can not be read.
409     */
410    private Document newDocument( final Page page, final String executionId ) throws IOException {
411        final String pageName = page.getName();
412        final List< DocumentAttribute > attrs = new ArrayList<>();
413        // These 2 are required as per
414        // https://docs.aws.amazon.com/kendra/latest/dg/data-source-custom.html#custom-required-attributes
415        attrs.add( newAttribute( "_data_source_id", dataSourceId ) );
416        attrs.add( newAttribute( "_data_source_sync_job_execution_id", executionId ) );
417
418        final String title = TextUtil.beautifyString( pageName );
419        ByteBuffer blob;
420        ContentType contentType = ContentType.PLAIN_TEXT;
421        if ( page instanceof final Attachment attachment ) {
422            InputStream is = null;
423            try {
424                final String filename = attachment.getFileName();
425                contentType = getContentType( filename );
426                is = engine.getManager( AttachmentManager.class ).getAttachmentStream( attachment );
427                blob = ByteBuffer.wrap( IOUtils.toByteArray( is ) );
428            } catch ( final ProviderException e ) {
429                throw new IOException( e );
430            } finally {
431                IOUtils.closeQuietly( is, null );
432            }
433            // contentType should be set to its real value
434        } else {
435            final String text = engine.getManager( PageManager.class ).getPureText( page );
436            blob = ByteBuffer.wrap( text.getBytes( StandardCharsets.UTF_8 ) );
437        }
438        return new Document().withId( pageName ).withTitle( title ).withAttributes( attrs ).withBlob( blob )
439                .withContentType( contentType );
440    }
441
442    private DocumentAttribute newAttribute( final String key, final String value ) {
443        return new DocumentAttribute().withKey( key ).withValue( new DocumentAttributeValue().withStringValue( value ) );
444    }
445
446    @SuppressWarnings( "unchecked" )
447    private ContentType getContentType( final String filename ) {
448        final String extention = FilenameUtils.getExtension( filename );
449        final Map< String, String > ct = ( Map< String, String > ) this.contentTypes.get( "ContentTypes" );
450        return ContentType.fromValue( ct.getOrDefault( extention, ContentType.PLAIN_TEXT.name() ) );
451    }
452
453    /**
454     * Updater thread that updates Kendra indexes.
455     */
456    private static final class KendraUpdater extends WikiBackgroundThread {
457        static final int INDEX_DELAY = 5;
458        static final int INITIAL_DELAY = 10;
459        private final KendraSearchProvider provider;
460
461        private final int initialDelay;
462
463        private WatchDog watchdog;
464
465        private KendraUpdater( final Engine engine, final KendraSearchProvider provider, final int initialDelay, final int indexDelay ) {
466            super( engine, indexDelay );
467            this.provider = provider;
468            this.initialDelay = initialDelay;
469            setName( "JSPWiki Kendra Indexer" );
470        }
471
472        @Override
473        public void startupTask() throws Exception {
474            watchdog = WatchDog.getCurrentWatchDog( getEngine() );
475            try {
476                Thread.sleep( initialDelay * 1000L );
477            } catch ( final InterruptedException e ) {
478                throw new InternalWikiException( "Interrupted while waiting to start.", e );
479            }
480            watchdog.enterState( "Full reindex" );
481            provider.initializeIndexAndDataSource();
482            provider.doFullReindex();
483            watchdog.exitState();
484        }
485
486        @Override
487        public void backgroundTask() {
488            watchdog.enterState( "Reindexing ...", 60 );
489            provider.doPartialReindex();
490            watchdog.exitState();
491        }
492    }
493
494    private static class SearchResultImpl implements SearchResult {
495
496        private final Page page;
497        private final int score;
498        private final String[] contexts;
499
500        public SearchResultImpl( final Page page, final int score, final String[] contexts ) {
501            this.page = page;
502            this.score = score;
503            this.contexts = contexts != null ? contexts.clone() : null;
504        }
505
506        @Override
507        public Page getPage() {
508            return this.page;
509        }
510
511        @Override
512        public int getScore() {
513            return this.score;
514        }
515
516        @Override
517        public String[] getContexts() {
518            return this.contexts;
519        }
520    }
521
522    public AWSkendra getKendra() {
523        return kendra;
524    }
525
526    public void setKendra( final AWSkendra kendra ) {
527        this.kendra = kendra;
528    }
529
530    protected AWSkendra buildClient() {
531        return AWSkendraClientBuilder.defaultClient();
532    }
533
534    public String getIndexName() {
535        return indexName;
536    }
537
538    public String getDataSourceName() {
539        return dataSourceName;
540    }
541
542}