Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -1215,18 +1215,11 @@ public Response search(@QueryParam("q") String queryTerm) {
HashSet<String> userAndRoles = new HashSet<>();
userAndRoles.add(principal);
userAndRoles.addAll(roles);
List<Map<String, String>> notesFound = noteSearchService.query(queryTerm);
for (int i = 0; i < notesFound.size(); i++) {
String[] ids = notesFound.get(i).get("id").split("/", 2);
String noteId = ids[0];
if (!authorizationService.isOwner(noteId, userAndRoles) &&
!authorizationService.isReader(noteId, userAndRoles) &&
!authorizationService.isWriter(noteId, userAndRoles) &&
!authorizationService.isRunner(noteId, userAndRoles)) {
notesFound.remove(i);
i--;
}
}
// isReader() already covers owners, writers and runners. Handing the check to the
// search service keeps it in front of the cutoff, so a caller with access to few notes
// is still served a full result set of the notes it may read.
List<Map<String, String>> notesFound = noteSearchService.query(queryTerm,
noteId -> authorizationService.isReader(noteId, userAndRoles));
LOGGER.info("{} notes found", notesFound.size());
return new JsonResponse<>(Status.OK, notesFound).build();
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.function.Predicate;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
Expand Down Expand Up @@ -509,15 +510,11 @@ private String buildParagraphText(String noteName, Paragraph p) {
// ---- SearchService implementation ----

@Override
// TODO(ZEPPELIN-6414): Accept user/roles (or a readability Predicate) and apply the auth
// filter before Phase-1 table collection and before the top-K cutoff. Currently the REST
// layer filters after truncation, which can hide results the caller is authorized for and
// lets inaccessible notes contaminate the table-boost ranking. Requires a SearchService
// interface change that also affects LuceneSearch.
public List<Map<String, String>> query(String queryStr) {
public List<Map<String, String>> query(String queryStr, Predicate<String> readable) {
if (StringUtils.isBlank(queryStr) || index.isEmpty()) {
return Collections.emptyList();
}
Map<String, Boolean> readableNotes = new HashMap<>();

float[] queryEmbedding = embed(queryStr);
String queryLower = queryStr.toLowerCase(Locale.ROOT);
Expand All @@ -527,6 +524,12 @@ public List<Map<String, String>> query(String queryStr) {
indexLock.readLock().lock();
try {
for (Map.Entry<String, IndexEntry> entry : index.entrySet()) {
// Dropping the entries here keeps them out of the table weights below and out of
// the cutoff, so the caller is served its own top results and not what is left of
// everyone's top results.
if (!readableNotes.computeIfAbsent(noteIdOf(entry.getKey()), readable::test)) {
continue;
}
float sim = cosineSimilarity(queryEmbedding, entry.getValue().embedding);
IndexEntry ie = entry.getValue();
if (ie.text != null && ie.text.toLowerCase(Locale.ROOT).contains(queryLower)) {
Expand Down Expand Up @@ -612,6 +615,16 @@ public List<Map<String, String>> query(String queryStr) {
return results;
}

/**
* The key of an indexed entry is either a noteId or a noteId followed by the paragraph.
*
* @see #formatId(String, Paragraph)
*/
private static String noteIdOf(String docId) {
int separator = docId.indexOf('/');
return separator < 0 ? docId : docId.substring(0, separator);
}

@Override
public void addNoteIndex(String noteId) {
try {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,10 +21,13 @@
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Collections;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.function.Predicate;
import javax.annotation.PreDestroy;
import jakarta.inject.Inject;

Expand Down Expand Up @@ -76,6 +79,12 @@ public class LuceneSearch extends SearchService {
private static final String SEARCH_FIELD_TITLE = "header";
private static final String PARAGRAPH = "paragraph";
private static final String ID_FIELD = "id";
/** Number of results a query returns at most. */
private static final int MAX_RESULTS = 20;
/** Number of hits pulled from the index per round while looking for readable ones. */
private static final int HIT_BATCH_SIZE = 20;
/** Only the id is needed to tell whether the caller may read a hit. */
private static final Set<String> ID_FIELD_ONLY = Collections.singleton(ID_FIELD);

private final Directory indexDirectory;
private final IndexWriter indexWriter;
Expand Down Expand Up @@ -113,7 +122,7 @@ public LuceneSearch(ZeppelinConfiguration zConf, Notebook notebook) throws IOExc
* @see org.apache.zeppelin.search.Search#query(java.lang.String)
*/
@Override
public List<Map<String, String>> query(String queryStr) {
public List<Map<String, String>> query(String queryStr, Predicate<String> readable) {
if (null == indexDirectory) {
throw new IllegalStateException(
"Something went wrong on instance creation time, index dir is null");
Expand All @@ -134,7 +143,7 @@ public List<Map<String, String>> query(String queryStr) {
SimpleHTMLFormatter htmlFormatter = new SimpleHTMLFormatter();
Highlighter highlighter = new Highlighter(htmlFormatter, new QueryScorer(query));

result = doSearch(indexSearcher, query, analyzer, highlighter);
result = doSearch(indexSearcher, query, analyzer, highlighter, readable);
} catch (IOException e) {
LOGGER.error("Failed to open index dir {}, make sure indexing finished OK", indexDirectory, e);
} catch (ParseException e) {
Expand All @@ -144,64 +153,40 @@ public List<Map<String, String>> query(String queryStr) {
}

private List<Map<String, String>> doSearch(
IndexSearcher searcher, Query query, Analyzer analyzer, Highlighter highlighter) {
IndexSearcher searcher, Query query, Analyzer analyzer, Highlighter highlighter,
Predicate<String> readable) {
List<Map<String, String>> matchingParagraphs = new ArrayList<>();
ScoreDoc[] hits;
Map<String, Boolean> readableNotes = new HashMap<>();
try {
hits = searcher.search(query, 20).scoreDocs;
for (int i = 0; i < hits.length; i++) {
LOGGER.debug("doc={} score={}", hits[i].doc, hits[i].score);

int id = hits[i].doc;
Document doc = searcher.doc(id);
String path = doc.get(ID_FIELD);
if (path != null) {
LOGGER.debug( "{}. {}", (i + 1), path);
String title = doc.get("title");
if (title != null) {
LOGGER.debug(" Title: {}", doc.get("title"));
// Walk the hits in score order and keep the ones the caller may read until the result
// set is full or the hits run out. Reading is checked here and not on the result set,
// because a cut that runs first would hide results the caller is allowed to see.
ScoreDoc lastHit = null;
while (matchingParagraphs.size() < MAX_RESULTS) {
ScoreDoc[] hits = (lastHit == null
? searcher.search(query, HIT_BATCH_SIZE)
: searcher.searchAfter(lastHit, query, HIT_BATCH_SIZE)).scoreDocs;
if (hits.length == 0) {
break;
}
lastHit = hits[hits.length - 1];
for (ScoreDoc hit : hits) {
if (matchingParagraphs.size() >= MAX_RESULTS) {
break;
}

String text = doc.get(SEARCH_FIELD_TEXT);
String header = doc.get(SEARCH_FIELD_TITLE);
String fragment = "";

if (text != null) {
TokenStream tokenStream =
TokenSources.getTokenStream(
searcher.getIndexReader(), id, SEARCH_FIELD_TEXT, analyzer);
TextFragment[] frags = highlighter.getBestTextFragments(tokenStream, text, true, 3);
LOGGER.debug(" {} fragments found for query '{}'", frags.length, query);
for (TextFragment frag : frags) {
if ((frag != null) && (frag.getScore() > 0)) {
LOGGER.debug(" Fragment: {}", frag);
}
}
fragment = (frags != null && frags.length > 0) ? frags[0].toString() : "";
LOGGER.debug("doc={} score={}", hit.doc, hit.score);
// Read the id alone to decide on a hit. The rest of the document is only worth
// loading for the hits that end up in the result set.
String path = searcher.doc(hit.doc, ID_FIELD_ONLY).get(ID_FIELD);
if (path == null) {
LOGGER.info("No {} for this document", ID_FIELD);
continue;
}

if (header != null) {
TokenStream tokenTitle =
TokenSources.getTokenStream(
searcher.getIndexReader(), id, SEARCH_FIELD_TITLE, analyzer);
TextFragment[] frgTitle = highlighter.getBestTextFragments(tokenTitle, header, true, 3);
header = (frgTitle != null && frgTitle.length > 0) ? frgTitle[0].toString() : "";
} else {
header = "";
if (!readableNotes.computeIfAbsent(noteIdOf(path), readable::test)) {
continue;
}
matchingParagraphs.add(
ImmutableMap.<String, String>builder()
.put("id", path)
.put("name", title)
.put("snippet", fragment)
.put("text", text)
.put("header", header)
.put("title", header)
.put("tables", "")
.put("output", "")
.build());
} else {
LOGGER.info("{}. No {} for this document", i + 1, ID_FIELD);
toMatch(searcher, query, analyzer, highlighter, hit.doc, searcher.doc(hit.doc)));
}
}
} catch (IOException | InvalidTokenOffsetsException e) {
Expand All @@ -210,6 +195,55 @@ private List<Map<String, String>> doSearch(
return matchingParagraphs;
}

private Map<String, String> toMatch(IndexSearcher searcher, Query query, Analyzer analyzer,
Highlighter highlighter, int docId, Document doc)
throws IOException, InvalidTokenOffsetsException {
String path = doc.get(ID_FIELD);
String title = doc.get("title");
String text = doc.get(SEARCH_FIELD_TEXT);
String header = doc.get(SEARCH_FIELD_TITLE);
String fragment = "";

if (text != null) {
TokenStream tokenStream =
TokenSources.getTokenStream(
searcher.getIndexReader(), docId, SEARCH_FIELD_TEXT, analyzer);
TextFragment[] frags = highlighter.getBestTextFragments(tokenStream, text, true, 3);
LOGGER.debug(" {} fragments found for query '{}'", frags.length, query);
fragment = (frags != null && frags.length > 0) ? frags[0].toString() : "";
}

if (header != null) {
TokenStream tokenTitle =
TokenSources.getTokenStream(
searcher.getIndexReader(), docId, SEARCH_FIELD_TITLE, analyzer);
TextFragment[] frgTitle = highlighter.getBestTextFragments(tokenTitle, header, true, 3);
header = (frgTitle != null && frgTitle.length > 0) ? frgTitle[0].toString() : "";
} else {
header = "";
}
return ImmutableMap.<String, String>builder()
.put("id", path)
.put("name", title)
.put("snippet", fragment)
.put("text", text)
.put("header", header)
.put("title", header)
.put("tables", "")
.put("output", "")
.build();
}

/**
* The id of an indexed document is either a noteId or a noteId followed by the paragraph.
*
* @see #formatId(String, Paragraph)
*/
private static String noteIdOf(String documentId) {
int separator = documentId.indexOf('/');
return separator < 0 ? documentId : documentId.substring(0, separator);
}

/* (non-Javadoc)
* @see org.apache.zeppelin.search.Search#updateIndexDoc(org.apache.zeppelin.notebook.Note)
*/
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.Predicate;

public class NoSearchService extends SearchService {

Expand All @@ -30,7 +31,7 @@ public NoSearchService() {
}

@Override
public List<Map<String, String>> query(String queryStr) {
public List<Map<String, String>> query(String queryStr, Predicate<String> readable) {
return Collections.emptyList();
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.function.Predicate;
import org.apache.zeppelin.notebook.NoteEventAsyncListener;
import javax.annotation.PreDestroy;

Expand All @@ -39,9 +40,12 @@ protected SearchService(String name) {
* Full-text search in all the notes
*
* @param queryStr a query
* @param readable tells for a noteId whether the caller may read it. Entries the caller
* cannot read are dropped before the result set is cut down, so that the
* caller is not served fewer results than it is allowed to see.
* @return A list of matching paragraphs (id, text, snippet w/ highlight)
*/
public abstract List<Map<String, String>> query(String queryStr);
public abstract List<Map<String, String>> query(String queryStr, Predicate<String> readable);

/**
* Updates note index for the given note, only update index of note meta info,
Expand Down
Loading
Loading