    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.StringUtils;
    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.*;
    058
    059import static java.lang.String.format;
    060
    061/**
    062 * Search provider that implements {link SearchProvider} using AWS Kendra for
    063 * indexing. Note that we are using a Custom DataSource which limits the
    064 * attributes that can be uploaded / searched for each page (as per
    065 * https://docs.aws.amazon.com/kendra/latest/dg/custom-attributes.html). This
    066 * could be overcome by using an S3 bucket where any custom attributes can be
    067 * added.
    068 *
    069 * @since 2.11.0
    070 */
    071public class KendraSearchProvider implements SearchProvider {
    072
    073    private static final Logger LOG = LogManager.getLogger( KendraSearchProvider.class );
    074    private Engine engine;
    075    private Properties properties;
    076    private Map< String, Object > contentTypes;
    077    private AWSkendra kendra;
    078    private String indexName;
    079    private String indexId;
    080    private String dataSourceName;
    081    private String dataSourceId;
    082
    083    private final List< Page > updates = Collections.synchronizedList( new ArrayList<>() );
    084
    085    private static final String PROP_KENDRA_INDEX_NAME = "jspwiki.kendra.indexName";
    086    private static final String PROP_KENDRA_DATA_SOURCE_NAME = "jspwiki.kendra.dataSourceName";
    087    private static final String PROP_KENDRA_INDEXDELAY = "jspwiki.kendra.indexdelay";
    088    private static final String PROP_KENDRA_INITIALDELAY = "jspwiki.kendra.initialdelay";
    089
    090    public KendraSearchProvider() {
    091    }
    092
    093    /**
    094     * {@inheritDoc}
    095     */
    096    @Override
    097    public void initialize( final Engine engine, final Properties properties ) throws NoRequiredPropertyException, IOException {
    098        this.engine = engine;
    099        this.properties = properties;
    100        this.contentTypes = getContentTypes();
    101
    102        setKendra( buildClient() );
    103
    104        this.indexName = TextUtil.getRequiredProperty( this.properties, PROP_KENDRA_INDEX_NAME );
    105        this.dataSourceName = TextUtil.getRequiredProperty( this.properties, PROP_KENDRA_DATA_SOURCE_NAME );
    106        final int initialDelay = TextUtil.getIntegerProperty( this.properties, PROP_KENDRA_INITIALDELAY, KendraUpdater.INITIAL_DELAY );
    107        final int indexDelay = TextUtil.getIntegerProperty( this.properties, PROP_KENDRA_INDEXDELAY, KendraUpdater.INDEX_DELAY );
    108
    109        // Start the Kendra update thread, which waits first for a little while
    110        // before starting to go through the "pages that need updating".
    111        if ( initialDelay >= 0 ) {
    112            final KendraUpdater updater = new KendraUpdater( engine, this, initialDelay, indexDelay );
    113            updater.start();
    114        }
    115    }
    116
    117    private Map< String, Object > getContentTypes() {
    118        final Gson gson = new GsonBuilder().create();
    119        try ( final InputStream in = KendraSearchProvider.class.getResourceAsStream( "content_types.json" ) ) {
    120            if ( in != null ) {
    121                final Type collectionType = new TypeToken< HashMap< String, Object > >() {
    122                }.getType();
    123                return gson.fromJson( new InputStreamReader( in ), collectionType );
    124            }
    125        } catch ( final IOException e ) {
    126            LOG.error( format( "Unable to load default propertyfile 'content_types.json': %s", e.getMessage() ), e );
    127        }
    128        return null;
    129    }
    130
    131    /**
    132     * {@inheritDoc}
    133     */
    134    @Override
    135    public String getProviderInfo() {
    136        return "KendraSearchProvider";
    137    }
    138
    139    /**
    140     * {@inheritDoc}
    141     */
    142    @Override
    143    public void pageRemoved( final Page page ) {
    144        final String pageName = page.getName();
    145        final BatchDeleteDocumentRequest request = new BatchDeleteDocumentRequest().withIndexId( indexId )
    146                .withDocumentIdList( pageName );
    147        final BatchDeleteDocumentResult result = getKendra().batchDeleteDocument( request );
    148        if (result.getFailedDocuments().isEmpty()) {
    149            LOG.debug( format( "Page '%s' was removed from index", pageName ) );
    150        } else {
    151            LOG.error( format( "Failed to remove Page '%s' from index", pageName ) );
    152        }
    153    }
    154
    155    /**
    156     * {@inheritDoc}
    157     */
    158    @Override
    159    public void reindexPage( final Page page ) {
    160        if ( page != null ) {
    161            updates.add( page );
    162            LOG.debug( format( "Scheduling page '%s' for indexing ...", page.getName() ) );
    163        }
    164    }
    165
    166    /**
    167     * {@inheritDoc}
    168     */
    169    @Override
    170    public Collection< SearchResult > findPages( final String query, final Context wikiContext ) throws ProviderException, IOException {
    171        final QueryRequest request = new QueryRequest().withIndexId( indexId ).withQueryText( query );
    172        final List< QueryResultItem > items;
    173        try {
    174            items = getKendra().query( request ).getResultItems();
    175        } catch ( final ThrottlingException e ) {
    176            LOG.error( "ThrottlingException. Skipping..." );
    177            return new ArrayList<>();
    178        }
    179        final List< SearchResult > searchResults = new ArrayList<>( items.size() );
    180        final AuthorizationManager mgr = engine.getManager( AuthorizationManager.class );
    181
    182        for ( final QueryResultItem item : items ) {
    183            switch( QueryResultType.fromValue( item.getType() ) ) {
    184                case DOCUMENT:
    185                    final String documentId = item.getDocumentId();
    186                    final String documentExcerpt = item.getDocumentExcerpt().getText();
    187                    final String scoreConfidence = item.getScoreAttributes().getScoreConfidence();
    188                    final Page page = this.engine.getManager( PageManager.class ).getPage( documentId, PageProvider.LATEST_VERSION );
    189                    if ( page != null ) {
    190                        final PagePermission pp = new PagePermission( page, PagePermission.VIEW_ACTION );
    191                        if ( mgr.checkPermission( wikiContext.getWikiSession(), pp ) ) {
    192                            final SearchResult searchResult = new SearchResultImpl( page, confidence2score( scoreConfidence ),
    193                                    new String[]{ documentExcerpt } );
    194                            searchResults.add( searchResult );
    195                        } else {
    196                            LOG.error( format( "Page '%s' is not accessible", documentId ) );
    197                        }
    198                    } else {
    199                        LOG.error(
    200                                format( "Kendra found a result page '%s' that could not be loaded, removing from index", documentId ) );
    201                        pageRemoved( Wiki.contents().page( this.engine, documentId ) );
    202                    }
    203                    break;
    204                default:
    205                    LOG.error( format( "Unknown query result type: %s", item.getType() ) );
    206            }
    207        }
    208        return searchResults;
    209    }
    210
    211    /**
    212     * This method initialize the AWS Kendra Index and Datasources to be used.
    213     */
    214    public void initializeIndexAndDataSource() {
    215        this.indexId = getIndexId( indexName );
    216        if ( this.indexId == null ) {
    217            final String message = format( "Index '%s' does not exist", indexName );
    218            LOG.error( message );
    219            throw new IllegalArgumentException( message );
    220        }
    221        this.dataSourceId = getDatasourceId( this.indexId, dataSourceName );
    222        if ( this.dataSourceId == null ) {
    223            final String message = format( "Datasource '%s' does not exist in index %s", dataSourceName, indexName );
    224            LOG.error( message );
    225            throw new IllegalArgumentException( message );
    226        }
    227    }
    228
    229    /**
    230     * Given an Kendra's Index name, returns the corresponding Index Id, or
    231     * {@code null} if it does not exists
    232     *
    233     * @param indexName the name of the index to look up
    234     * @return the index id or {@code null}
    235     */
    236    private String getIndexId( final String indexName ) {
    237        ListIndicesRequest request = new ListIndicesRequest();
    238        ListIndicesResult result = getKendra().listIndices( request );
    239        String nextToken = "";
    240        while ( nextToken != null ) {
    241            final List< IndexConfigurationSummary > items = result.getIndexConfigurationSummaryItems();
    242            if ( items == null || items.isEmpty() ) {
    243                return null;
    244            }
    245            for ( final IndexConfigurationSummary item : items ) {
    246                if ( StringUtils.equals( item.getName(), indexName ) ) {
    247                    return item.getId();
    248                }
    249            }
    250            nextToken = result.getNextToken();
    251            request = new ListIndicesRequest().withNextToken( result.getNextToken() );
    252            result = getKendra().listIndices( request );
    253        }
    254        return null;
    255    }
    256
    257    /**
    258     * Given an Kendra's Datasource name, returns the corresponding Datasource Id,
    259     * or {@code null} if it does not exists
    260     *
    261     * @param dataSourceName the name of the datasource to look up
    262     * @return the datasource id or {@code null}
    263     */
    264    private String getDatasourceId( final String indexId, final String dataSourceName ) {
    265        ListDataSourcesRequest request = new ListDataSourcesRequest().withIndexId( indexId );
    266        ListDataSourcesResult result = getKendra().listDataSources( request );
    267        String nextToken = "";
    268        while ( nextToken != null ) {
    269            final List< DataSourceSummary > items = result.getSummaryItems();
    270            if ( items == null || items.isEmpty() ) {
    271                return null;
    272            }
    273
    274            for ( final DataSourceSummary item : items ) {
    275                if ( StringUtils.equals( item.getName(), dataSourceName ) ) {
    276                    return item.getId();
    277                }
    278            }
    279            nextToken = result.getNextToken();
    280            request = new ListDataSourcesRequest().withNextToken( result.getNextToken() );
    281            result = getKendra().listDataSources( request );
    282        }
    283        return null;
    284    }
    285
    286    /*
    287     * Converts a SCORE Confidence from Kendra to an "equivalent" integer score
    288     */
    289    private int confidence2score( final String scoreConfidence ) {
    290        switch ( ScoreConfidence.fromValue( scoreConfidence ) ) {
    291            case VERY_HIGH:
    292                return 100;
    293            case HIGH:
    294                return 75;
    295            case MEDIUM:
    296                return 50;
    297            case LOW:
    298                return 25;
    299            default:
    300                return 0;
    301        }
    302    }
    303
    304    /**
    305     * This method re-index all the pages found in the Wiki. It is mainly used at
    306     * startup.
    307     *
    308     * @throws IOException in case some page can not be read
    309     */
    310    private void doFullReindex() throws IOException {
    311        try {
    312            final Collection< Page > pages = engine.getManager( PageManager.class ).getAllPages();
    313            if ( pages.isEmpty() ) {
    314                return;
    315            }
    316            LOG.debug( format( "Indexing all %d pages. Please wait ...", pages.size() ) );
    317            final String executionId = startExecution();
    318            for ( final Page page : pages ) {
    319                // Since I do not want to handle the size limit
    320                // (https://docs.aws.amazon.com/goto/WebAPI/kendra-2019-02-03/BatchPutDocument)
    321                // uploading documents one at a time
    322                indexOnePage( page, executionId );
    323            }
    324        } catch ( final ProviderException e ) {
    325            LOG.error( e.getMessage() );
    326            throw new IOException( e );
    327        } finally {
    328            stopExecution();
    329        }
    330    }
    331
    332    /**
    333     * This method re-index all pages marked as updated. It is used to periodically
    334     * index pages that have been modified
    335     */
    336    private void doPartialReindex() {
    337        if ( updates.isEmpty() ) {
    338            return;
    339        }
    340        LOG.debug( "Indexing updated pages. Please wait ..." );
    341        final String executionId = startExecution();
    342        synchronized ( updates ) {
    343            try {
    344                while (!updates.isEmpty()) {
    345                    indexOnePage( updates.remove( 0 ), executionId );
    346                }
    347            } finally {
    348                stopExecution();
    349            }
    350        }
    351    }
    352
    353    /**
    354     * Returns an ExecutiuonId that is required to keep track of the modifed
    355     * documents
    356     *
    357     * @return The execution id
    358     */
    359    private String startExecution() {
    360        final StartDataSourceSyncJobRequest request = new StartDataSourceSyncJobRequest().withIndexId( indexId )
    361                .withId( dataSourceId );
    362        final StartDataSourceSyncJobResult result = getKendra().startDataSourceSyncJob( request );
    363        return result.getExecutionId();
    364    }
    365
    366    /**
    367     * Stop the execution for the given index Id and DataSource Id.
    368     */
    369    private void stopExecution() {
    370        final StopDataSourceSyncJobRequest request = new StopDataSourceSyncJobRequest().withIndexId( indexId ).withId( dataSourceId );
    371        getKendra().stopDataSourceSyncJob( request );
    372    }
    373
    374    /**
    375     * Index on single {@link Page} into the Kendra Index
    376     *
    377     * @param page        the {@link Page} to index
    378     * @param executionId The Execution Id
    379     */
    380    private void indexOnePage( final Page page, final String executionId ) {
    381        final String pageName = page.getName();
    382        try {
    383            final Document document = newDocument( page, executionId );
    384            final BatchPutDocumentRequest request = new BatchPutDocumentRequest().withIndexId( indexId )
    385                    .withDocuments( document );
    386            final BatchPutDocumentResult result = getKendra().batchPutDocument( request );
    387            if (result.getFailedDocuments().isEmpty()) {
    388                LOG.info( format( "Successfully indexed Page '%s' as %s", page.getName(), document.getContentType() ) );
    389            } else {
    390                for ( final BatchPutDocumentResponseFailedDocument failedDocument : result.getFailedDocuments() ) {
    391                    LOG.error( format( "Failed to index Page '%s': %s", failedDocument.getId(), failedDocument.getErrorMessage() ) );
    392                }
    393            }
    394        } catch ( final IOException e ) {
    395            LOG.error( format( "Failed to index Page '%s': %s", pageName, e.getMessage() ) );
    396        }
    397    }
    398
    399
    400    /**
    401     * Given a {@link Page}, returns the corresponding Kendra {@link Document}.
    402     *
    403     * @param page        the {@link Page} to be indexed
    404     * @param executionId an execution id to identify when the {@link Page} was
    405     *                    indexed for the last time.
    406     * @return a {@link Document} containing the searchable attributes.
    407     * @throws IOException if the {@link Page}'s {@link Attachment} can not be read.
    408     */
    409    private Document newDocument( final Page page, final String executionId ) throws IOException {
    410        final String pageName = page.getName();
    411        final List< DocumentAttribute > attrs = new ArrayList<>();
    412        // These 2 are required as per
    413        // https://docs.aws.amazon.com/kendra/latest/dg/data-source-custom.html#custom-required-attributes
    414        attrs.add( newAttribute( "_data_source_id", dataSourceId ) );
    415        attrs.add( newAttribute( "_data_source_sync_job_execution_id", executionId ) );
    416
    417        final String title = TextUtil.beautifyString( pageName );
    418        ByteBuffer blob;
    419        ContentType contentType = ContentType.PLAIN_TEXT;
    420        if ( page instanceof Attachment ) {
    421            final Attachment attachment = ( Attachment ) page;
    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}