    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}