commands.py 36 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117
  1. import itertools
  2. import time
  3. from typing import Dict, List, Optional, Union
  4. from redis.client import Pipeline
  5. from redis.utils import deprecated_function
  6. from ..helpers import get_protocol_version, parse_to_dict
  7. from ._util import to_string
  8. from .aggregation import AggregateRequest, AggregateResult, Cursor
  9. from .document import Document
  10. from .query import Query
  11. from .result import Result
  12. from .suggestion import SuggestionParser
  13. NUMERIC = "NUMERIC"
  14. CREATE_CMD = "FT.CREATE"
  15. ALTER_CMD = "FT.ALTER"
  16. SEARCH_CMD = "FT.SEARCH"
  17. ADD_CMD = "FT.ADD"
  18. ADDHASH_CMD = "FT.ADDHASH"
  19. DROP_CMD = "FT.DROP"
  20. DROPINDEX_CMD = "FT.DROPINDEX"
  21. EXPLAIN_CMD = "FT.EXPLAIN"
  22. EXPLAINCLI_CMD = "FT.EXPLAINCLI"
  23. DEL_CMD = "FT.DEL"
  24. AGGREGATE_CMD = "FT.AGGREGATE"
  25. PROFILE_CMD = "FT.PROFILE"
  26. CURSOR_CMD = "FT.CURSOR"
  27. SPELLCHECK_CMD = "FT.SPELLCHECK"
  28. DICT_ADD_CMD = "FT.DICTADD"
  29. DICT_DEL_CMD = "FT.DICTDEL"
  30. DICT_DUMP_CMD = "FT.DICTDUMP"
  31. GET_CMD = "FT.GET"
  32. MGET_CMD = "FT.MGET"
  33. CONFIG_CMD = "FT.CONFIG"
  34. TAGVALS_CMD = "FT.TAGVALS"
  35. ALIAS_ADD_CMD = "FT.ALIASADD"
  36. ALIAS_UPDATE_CMD = "FT.ALIASUPDATE"
  37. ALIAS_DEL_CMD = "FT.ALIASDEL"
  38. INFO_CMD = "FT.INFO"
  39. SUGADD_COMMAND = "FT.SUGADD"
  40. SUGDEL_COMMAND = "FT.SUGDEL"
  41. SUGLEN_COMMAND = "FT.SUGLEN"
  42. SUGGET_COMMAND = "FT.SUGGET"
  43. SYNUPDATE_CMD = "FT.SYNUPDATE"
  44. SYNDUMP_CMD = "FT.SYNDUMP"
  45. NOOFFSETS = "NOOFFSETS"
  46. NOFIELDS = "NOFIELDS"
  47. NOHL = "NOHL"
  48. NOFREQS = "NOFREQS"
  49. MAXTEXTFIELDS = "MAXTEXTFIELDS"
  50. TEMPORARY = "TEMPORARY"
  51. STOPWORDS = "STOPWORDS"
  52. SKIPINITIALSCAN = "SKIPINITIALSCAN"
  53. WITHSCORES = "WITHSCORES"
  54. FUZZY = "FUZZY"
  55. WITHPAYLOADS = "WITHPAYLOADS"
  56. class SearchCommands:
  57. """Search commands."""
  58. def _parse_results(self, cmd, res, **kwargs):
  59. if get_protocol_version(self.client) in ["3", 3]:
  60. return res
  61. else:
  62. return self._RESP2_MODULE_CALLBACKS[cmd](res, **kwargs)
  63. def _parse_info(self, res, **kwargs):
  64. it = map(to_string, res)
  65. return dict(zip(it, it))
  66. def _parse_search(self, res, **kwargs):
  67. return Result(
  68. res,
  69. not kwargs["query"]._no_content,
  70. duration=kwargs["duration"],
  71. has_payload=kwargs["query"]._with_payloads,
  72. with_scores=kwargs["query"]._with_scores,
  73. )
  74. def _parse_aggregate(self, res, **kwargs):
  75. return self._get_aggregate_result(res, kwargs["query"], kwargs["has_cursor"])
  76. def _parse_profile(self, res, **kwargs):
  77. query = kwargs["query"]
  78. if isinstance(query, AggregateRequest):
  79. result = self._get_aggregate_result(res[0], query, query._cursor)
  80. else:
  81. result = Result(
  82. res[0],
  83. not query._no_content,
  84. duration=kwargs["duration"],
  85. has_payload=query._with_payloads,
  86. with_scores=query._with_scores,
  87. )
  88. return result, parse_to_dict(res[1])
  89. def _parse_spellcheck(self, res, **kwargs):
  90. corrections = {}
  91. if res == 0:
  92. return corrections
  93. for _correction in res:
  94. if isinstance(_correction, int) and _correction == 0:
  95. continue
  96. if len(_correction) != 3:
  97. continue
  98. if not _correction[2]:
  99. continue
  100. if not _correction[2][0]:
  101. continue
  102. # For spellcheck output
  103. # 1) 1) "TERM"
  104. # 2) "{term1}"
  105. # 3) 1) 1) "{score1}"
  106. # 2) "{suggestion1}"
  107. # 2) 1) "{score2}"
  108. # 2) "{suggestion2}"
  109. #
  110. # Following dictionary will be made
  111. # corrections = {
  112. # '{term1}': [
  113. # {'score': '{score1}', 'suggestion': '{suggestion1}'},
  114. # {'score': '{score2}', 'suggestion': '{suggestion2}'}
  115. # ]
  116. # }
  117. corrections[_correction[1]] = [
  118. {"score": _item[0], "suggestion": _item[1]} for _item in _correction[2]
  119. ]
  120. return corrections
  121. def _parse_config_get(self, res, **kwargs):
  122. return {kvs[0]: kvs[1] for kvs in res} if res else {}
  123. def _parse_syndump(self, res, **kwargs):
  124. return {res[i]: res[i + 1] for i in range(0, len(res), 2)}
  125. def batch_indexer(self, chunk_size=100):
  126. """
  127. Create a new batch indexer from the client with a given chunk size
  128. """
  129. return self.BatchIndexer(self, chunk_size=chunk_size)
  130. def create_index(
  131. self,
  132. fields,
  133. no_term_offsets=False,
  134. no_field_flags=False,
  135. stopwords=None,
  136. definition=None,
  137. max_text_fields=False,
  138. temporary=None,
  139. no_highlight=False,
  140. no_term_frequencies=False,
  141. skip_initial_scan=False,
  142. ):
  143. """
  144. Create the search index. The index must not already exist.
  145. ### Parameters:
  146. - **fields**: a list of TextField or NumericField objects
  147. - **no_term_offsets**: If true, we will not save term offsets in
  148. the index
  149. - **no_field_flags**: If true, we will not save field flags that
  150. allow searching in specific fields
  151. - **stopwords**: If not None, we create the index with this custom
  152. stopword list. The list can be empty
  153. - **max_text_fields**: If true, we will encode indexes as if there
  154. were more than 32 text fields which allows you to add additional
  155. fields (beyond 32).
  156. - **temporary**: Create a lightweight temporary index which will
  157. expire after the specified period of inactivity (in seconds). The
  158. internal idle timer is reset whenever the index is searched or added to.
  159. - **no_highlight**: If true, disabling highlighting support.
  160. Also implied by no_term_offsets.
  161. - **no_term_frequencies**: If true, we avoid saving the term frequencies
  162. in the index.
  163. - **skip_initial_scan**: If true, we do not scan and index.
  164. For more information see `FT.CREATE <https://redis.io/commands/ft.create>`_.
  165. """ # noqa
  166. args = [CREATE_CMD, self.index_name]
  167. if definition is not None:
  168. args += definition.args
  169. if max_text_fields:
  170. args.append(MAXTEXTFIELDS)
  171. if temporary is not None and isinstance(temporary, int):
  172. args.append(TEMPORARY)
  173. args.append(temporary)
  174. if no_term_offsets:
  175. args.append(NOOFFSETS)
  176. if no_highlight:
  177. args.append(NOHL)
  178. if no_field_flags:
  179. args.append(NOFIELDS)
  180. if no_term_frequencies:
  181. args.append(NOFREQS)
  182. if skip_initial_scan:
  183. args.append(SKIPINITIALSCAN)
  184. if stopwords is not None and isinstance(stopwords, (list, tuple, set)):
  185. args += [STOPWORDS, len(stopwords)]
  186. if len(stopwords) > 0:
  187. args += list(stopwords)
  188. args.append("SCHEMA")
  189. try:
  190. args += list(itertools.chain(*(f.redis_args() for f in fields)))
  191. except TypeError:
  192. args += fields.redis_args()
  193. return self.execute_command(*args)
  194. def alter_schema_add(self, fields: List[str]):
  195. """
  196. Alter the existing search index by adding new fields. The index
  197. must already exist.
  198. ### Parameters:
  199. - **fields**: a list of Field objects to add for the index
  200. For more information see `FT.ALTER <https://redis.io/commands/ft.alter>`_.
  201. """ # noqa
  202. args = [ALTER_CMD, self.index_name, "SCHEMA", "ADD"]
  203. try:
  204. args += list(itertools.chain(*(f.redis_args() for f in fields)))
  205. except TypeError:
  206. args += fields.redis_args()
  207. return self.execute_command(*args)
  208. def dropindex(self, delete_documents: bool = False):
  209. """
  210. Drop the index if it exists.
  211. Replaced `drop_index` in RediSearch 2.0.
  212. Default behavior was changed to not delete the indexed documents.
  213. ### Parameters:
  214. - **delete_documents**: If `True`, all documents will be deleted.
  215. For more information see `FT.DROPINDEX <https://redis.io/commands/ft.dropindex>`_.
  216. """ # noqa
  217. delete_str = "DD" if delete_documents else ""
  218. return self.execute_command(DROPINDEX_CMD, self.index_name, delete_str)
  219. def _add_document(
  220. self,
  221. doc_id,
  222. conn=None,
  223. nosave=False,
  224. score=1.0,
  225. payload=None,
  226. replace=False,
  227. partial=False,
  228. language=None,
  229. no_create=False,
  230. **fields,
  231. ):
  232. """
  233. Internal add_document used for both batch and single doc indexing
  234. """
  235. if partial or no_create:
  236. replace = True
  237. args = [ADD_CMD, self.index_name, doc_id, score]
  238. if nosave:
  239. args.append("NOSAVE")
  240. if payload is not None:
  241. args.append("PAYLOAD")
  242. args.append(payload)
  243. if replace:
  244. args.append("REPLACE")
  245. if partial:
  246. args.append("PARTIAL")
  247. if no_create:
  248. args.append("NOCREATE")
  249. if language:
  250. args += ["LANGUAGE", language]
  251. args.append("FIELDS")
  252. args += list(itertools.chain(*fields.items()))
  253. if conn is not None:
  254. return conn.execute_command(*args)
  255. return self.execute_command(*args)
  256. def _add_document_hash(
  257. self, doc_id, conn=None, score=1.0, language=None, replace=False
  258. ):
  259. """
  260. Internal add_document_hash used for both batch and single doc indexing
  261. """
  262. args = [ADDHASH_CMD, self.index_name, doc_id, score]
  263. if replace:
  264. args.append("REPLACE")
  265. if language:
  266. args += ["LANGUAGE", language]
  267. if conn is not None:
  268. return conn.execute_command(*args)
  269. return self.execute_command(*args)
  270. @deprecated_function(
  271. version="2.0.0", reason="deprecated since redisearch 2.0, call hset instead"
  272. )
  273. def add_document(
  274. self,
  275. doc_id: str,
  276. nosave: bool = False,
  277. score: float = 1.0,
  278. payload: bool = None,
  279. replace: bool = False,
  280. partial: bool = False,
  281. language: Optional[str] = None,
  282. no_create: str = False,
  283. **fields: List[str],
  284. ):
  285. """
  286. Add a single document to the index.
  287. ### Parameters
  288. - **doc_id**: the id of the saved document.
  289. - **nosave**: if set to true, we just index the document, and don't
  290. save a copy of it. This means that searches will just
  291. return ids.
  292. - **score**: the document ranking, between 0.0 and 1.0
  293. - **payload**: optional inner-index payload we can save for fast
  294. i access in scoring functions
  295. - **replace**: if True, and the document already is in the index,
  296. we perform an update and reindex the document
  297. - **partial**: if True, the fields specified will be added to the
  298. existing document.
  299. This has the added benefit that any fields specified
  300. with `no_index`
  301. will not be reindexed again. Implies `replace`
  302. - **language**: Specify the language used for document tokenization.
  303. - **no_create**: if True, the document is only updated and reindexed
  304. if it already exists.
  305. If the document does not exist, an error will be
  306. returned. Implies `replace`
  307. - **fields** kwargs dictionary of the document fields to be saved
  308. and/or indexed.
  309. NOTE: Geo points shoule be encoded as strings of "lon,lat"
  310. """ # noqa
  311. return self._add_document(
  312. doc_id,
  313. conn=None,
  314. nosave=nosave,
  315. score=score,
  316. payload=payload,
  317. replace=replace,
  318. partial=partial,
  319. language=language,
  320. no_create=no_create,
  321. **fields,
  322. )
  323. @deprecated_function(
  324. version="2.0.0", reason="deprecated since redisearch 2.0, call hset instead"
  325. )
  326. def add_document_hash(self, doc_id, score=1.0, language=None, replace=False):
  327. """
  328. Add a hash document to the index.
  329. ### Parameters
  330. - **doc_id**: the document's id. This has to be an existing HASH key
  331. in Redis that will hold the fields the index needs.
  332. - **score**: the document ranking, between 0.0 and 1.0
  333. - **replace**: if True, and the document already is in the index, we
  334. perform an update and reindex the document
  335. - **language**: Specify the language used for document tokenization.
  336. """ # noqa
  337. return self._add_document_hash(
  338. doc_id, conn=None, score=score, language=language, replace=replace
  339. )
  340. def delete_document(self, doc_id, conn=None, delete_actual_document=False):
  341. """
  342. Delete a document from index
  343. Returns 1 if the document was deleted, 0 if not
  344. ### Parameters
  345. - **delete_actual_document**: if set to True, RediSearch also delete
  346. the actual document if it is in the index
  347. """ # noqa
  348. args = [DEL_CMD, self.index_name, doc_id]
  349. if delete_actual_document:
  350. args.append("DD")
  351. if conn is not None:
  352. return conn.execute_command(*args)
  353. return self.execute_command(*args)
  354. def load_document(self, id):
  355. """
  356. Load a single document by id
  357. """
  358. fields = self.client.hgetall(id)
  359. f2 = {to_string(k): to_string(v) for k, v in fields.items()}
  360. fields = f2
  361. try:
  362. del fields["id"]
  363. except KeyError:
  364. pass
  365. return Document(id=id, **fields)
  366. def get(self, *ids):
  367. """
  368. Returns the full contents of multiple documents.
  369. ### Parameters
  370. - **ids**: the ids of the saved documents.
  371. """
  372. return self.execute_command(MGET_CMD, self.index_name, *ids)
  373. def info(self):
  374. """
  375. Get info an stats about the the current index, including the number of
  376. documents, memory consumption, etc
  377. For more information see `FT.INFO <https://redis.io/commands/ft.info>`_.
  378. """
  379. res = self.execute_command(INFO_CMD, self.index_name)
  380. return self._parse_results(INFO_CMD, res)
  381. def get_params_args(
  382. self, query_params: Union[Dict[str, Union[str, int, float, bytes]], None]
  383. ):
  384. if query_params is None:
  385. return []
  386. args = []
  387. if len(query_params) > 0:
  388. args.append("params")
  389. args.append(len(query_params) * 2)
  390. for key, value in query_params.items():
  391. args.append(key)
  392. args.append(value)
  393. return args
  394. def _mk_query_args(
  395. self, query, query_params: Union[Dict[str, Union[str, int, float, bytes]], None]
  396. ):
  397. args = [self.index_name]
  398. if isinstance(query, str):
  399. # convert the query from a text to a query object
  400. query = Query(query)
  401. if not isinstance(query, Query):
  402. raise ValueError(f"Bad query type {type(query)}")
  403. args += query.get_args()
  404. args += self.get_params_args(query_params)
  405. return args, query
  406. def search(
  407. self,
  408. query: Union[str, Query],
  409. query_params: Union[Dict[str, Union[str, int, float, bytes]], None] = None,
  410. ):
  411. """
  412. Search the index for a given query, and return a result of documents
  413. ### Parameters
  414. - **query**: the search query. Either a text for simple queries with
  415. default parameters, or a Query object for complex queries.
  416. See RediSearch's documentation on query format
  417. For more information see `FT.SEARCH <https://redis.io/commands/ft.search>`_.
  418. """ # noqa
  419. args, query = self._mk_query_args(query, query_params=query_params)
  420. st = time.time()
  421. res = self.execute_command(SEARCH_CMD, *args)
  422. if isinstance(res, Pipeline):
  423. return res
  424. return self._parse_results(
  425. SEARCH_CMD, res, query=query, duration=(time.time() - st) * 1000.0
  426. )
  427. def explain(
  428. self,
  429. query: Union[str, Query],
  430. query_params: Dict[str, Union[str, int, float]] = None,
  431. ):
  432. """Returns the execution plan for a complex query.
  433. For more information see `FT.EXPLAIN <https://redis.io/commands/ft.explain>`_.
  434. """ # noqa
  435. args, query_text = self._mk_query_args(query, query_params=query_params)
  436. return self.execute_command(EXPLAIN_CMD, *args)
  437. def explain_cli(self, query: Union[str, Query]): # noqa
  438. raise NotImplementedError("EXPLAINCLI will not be implemented.")
  439. def aggregate(
  440. self,
  441. query: Union[str, Query],
  442. query_params: Dict[str, Union[str, int, float]] = None,
  443. ):
  444. """
  445. Issue an aggregation query.
  446. ### Parameters
  447. **query**: This can be either an `AggregateRequest`, or a `Cursor`
  448. An `AggregateResult` object is returned. You can access the rows from
  449. its `rows` property, which will always yield the rows of the result.
  450. For more information see `FT.AGGREGATE <https://redis.io/commands/ft.aggregate>`_.
  451. """ # noqa
  452. if isinstance(query, AggregateRequest):
  453. has_cursor = bool(query._cursor)
  454. cmd = [AGGREGATE_CMD, self.index_name] + query.build_args()
  455. elif isinstance(query, Cursor):
  456. has_cursor = True
  457. cmd = [CURSOR_CMD, "READ", self.index_name] + query.build_args()
  458. else:
  459. raise ValueError("Bad query", query)
  460. cmd += self.get_params_args(query_params)
  461. raw = self.execute_command(*cmd)
  462. return self._parse_results(
  463. AGGREGATE_CMD, raw, query=query, has_cursor=has_cursor
  464. )
  465. def _get_aggregate_result(
  466. self, raw: List, query: Union[str, Query, AggregateRequest], has_cursor: bool
  467. ):
  468. if has_cursor:
  469. if isinstance(query, Cursor):
  470. query.cid = raw[1]
  471. cursor = query
  472. else:
  473. cursor = Cursor(raw[1])
  474. raw = raw[0]
  475. else:
  476. cursor = None
  477. if isinstance(query, AggregateRequest) and query._with_schema:
  478. schema = raw[0]
  479. rows = raw[2:]
  480. else:
  481. schema = None
  482. rows = raw[1:]
  483. return AggregateResult(rows, cursor, schema)
  484. def profile(
  485. self,
  486. query: Union[str, Query, AggregateRequest],
  487. limited: bool = False,
  488. query_params: Optional[Dict[str, Union[str, int, float]]] = None,
  489. ):
  490. """
  491. Performs a search or aggregate command and collects performance
  492. information.
  493. ### Parameters
  494. **query**: This can be either an `AggregateRequest`, `Query` or string.
  495. **limited**: If set to True, removes details of reader iterator.
  496. **query_params**: Define one or more value parameters.
  497. Each parameter has a name and a value.
  498. """
  499. st = time.time()
  500. cmd = [PROFILE_CMD, self.index_name, ""]
  501. if limited:
  502. cmd.append("LIMITED")
  503. cmd.append("QUERY")
  504. if isinstance(query, AggregateRequest):
  505. cmd[2] = "AGGREGATE"
  506. cmd += query.build_args()
  507. elif isinstance(query, Query):
  508. cmd[2] = "SEARCH"
  509. cmd += query.get_args()
  510. cmd += self.get_params_args(query_params)
  511. else:
  512. raise ValueError("Must provide AggregateRequest object or Query object.")
  513. res = self.execute_command(*cmd)
  514. return self._parse_results(
  515. PROFILE_CMD, res, query=query, duration=(time.time() - st) * 1000.0
  516. )
  517. def spellcheck(self, query, distance=None, include=None, exclude=None):
  518. """
  519. Issue a spellcheck query
  520. ### Parameters
  521. **query**: search query.
  522. **distance***: the maximal Levenshtein distance for spelling
  523. suggestions (default: 1, max: 4).
  524. **include**: specifies an inclusion custom dictionary.
  525. **exclude**: specifies an exclusion custom dictionary.
  526. For more information see `FT.SPELLCHECK <https://redis.io/commands/ft.spellcheck>`_.
  527. """ # noqa
  528. cmd = [SPELLCHECK_CMD, self.index_name, query]
  529. if distance:
  530. cmd.extend(["DISTANCE", distance])
  531. if include:
  532. cmd.extend(["TERMS", "INCLUDE", include])
  533. if exclude:
  534. cmd.extend(["TERMS", "EXCLUDE", exclude])
  535. res = self.execute_command(*cmd)
  536. return self._parse_results(SPELLCHECK_CMD, res)
  537. def dict_add(self, name: str, *terms: List[str]):
  538. """Adds terms to a dictionary.
  539. ### Parameters
  540. - **name**: Dictionary name.
  541. - **terms**: List of items for adding to the dictionary.
  542. For more information see `FT.DICTADD <https://redis.io/commands/ft.dictadd>`_.
  543. """ # noqa
  544. cmd = [DICT_ADD_CMD, name]
  545. cmd.extend(terms)
  546. return self.execute_command(*cmd)
  547. def dict_del(self, name: str, *terms: List[str]):
  548. """Deletes terms from a dictionary.
  549. ### Parameters
  550. - **name**: Dictionary name.
  551. - **terms**: List of items for removing from the dictionary.
  552. For more information see `FT.DICTDEL <https://redis.io/commands/ft.dictdel>`_.
  553. """ # noqa
  554. cmd = [DICT_DEL_CMD, name]
  555. cmd.extend(terms)
  556. return self.execute_command(*cmd)
  557. def dict_dump(self, name: str):
  558. """Dumps all terms in the given dictionary.
  559. ### Parameters
  560. - **name**: Dictionary name.
  561. For more information see `FT.DICTDUMP <https://redis.io/commands/ft.dictdump>`_.
  562. """ # noqa
  563. cmd = [DICT_DUMP_CMD, name]
  564. return self.execute_command(*cmd)
  565. def config_set(self, option: str, value: str) -> bool:
  566. """Set runtime configuration option.
  567. ### Parameters
  568. - **option**: the name of the configuration option.
  569. - **value**: a value for the configuration option.
  570. For more information see `FT.CONFIG SET <https://redis.io/commands/ft.config-set>`_.
  571. """ # noqa
  572. cmd = [CONFIG_CMD, "SET", option, value]
  573. raw = self.execute_command(*cmd)
  574. return raw == "OK"
  575. def config_get(self, option: str) -> str:
  576. """Get runtime configuration option value.
  577. ### Parameters
  578. - **option**: the name of the configuration option.
  579. For more information see `FT.CONFIG GET <https://redis.io/commands/ft.config-get>`_.
  580. """ # noqa
  581. cmd = [CONFIG_CMD, "GET", option]
  582. res = self.execute_command(*cmd)
  583. return self._parse_results(CONFIG_CMD, res)
  584. def tagvals(self, tagfield: str):
  585. """
  586. Return a list of all possible tag values
  587. ### Parameters
  588. - **tagfield**: Tag field name
  589. For more information see `FT.TAGVALS <https://redis.io/commands/ft.tagvals>`_.
  590. """ # noqa
  591. return self.execute_command(TAGVALS_CMD, self.index_name, tagfield)
  592. def aliasadd(self, alias: str):
  593. """
  594. Alias a search index - will fail if alias already exists
  595. ### Parameters
  596. - **alias**: Name of the alias to create
  597. For more information see `FT.ALIASADD <https://redis.io/commands/ft.aliasadd>`_.
  598. """ # noqa
  599. return self.execute_command(ALIAS_ADD_CMD, alias, self.index_name)
  600. def aliasupdate(self, alias: str):
  601. """
  602. Updates an alias - will fail if alias does not already exist
  603. ### Parameters
  604. - **alias**: Name of the alias to create
  605. For more information see `FT.ALIASUPDATE <https://redis.io/commands/ft.aliasupdate>`_.
  606. """ # noqa
  607. return self.execute_command(ALIAS_UPDATE_CMD, alias, self.index_name)
  608. def aliasdel(self, alias: str):
  609. """
  610. Removes an alias to a search index
  611. ### Parameters
  612. - **alias**: Name of the alias to delete
  613. For more information see `FT.ALIASDEL <https://redis.io/commands/ft.aliasdel>`_.
  614. """ # noqa
  615. return self.execute_command(ALIAS_DEL_CMD, alias)
  616. def sugadd(self, key, *suggestions, **kwargs):
  617. """
  618. Add suggestion terms to the AutoCompleter engine. Each suggestion has
  619. a score and string.
  620. If kwargs["increment"] is true and the terms are already in the
  621. server's dictionary, we increment their scores.
  622. For more information see `FT.SUGADD <https://redis.io/commands/ft.sugadd/>`_.
  623. """ # noqa
  624. # If Transaction is not False it will MULTI/EXEC which will error
  625. pipe = self.pipeline(transaction=False)
  626. for sug in suggestions:
  627. args = [SUGADD_COMMAND, key, sug.string, sug.score]
  628. if kwargs.get("increment"):
  629. args.append("INCR")
  630. if sug.payload:
  631. args.append("PAYLOAD")
  632. args.append(sug.payload)
  633. pipe.execute_command(*args)
  634. return pipe.execute()[-1]
  635. def suglen(self, key: str) -> int:
  636. """
  637. Return the number of entries in the AutoCompleter index.
  638. For more information see `FT.SUGLEN <https://redis.io/commands/ft.suglen>`_.
  639. """ # noqa
  640. return self.execute_command(SUGLEN_COMMAND, key)
  641. def sugdel(self, key: str, string: str) -> int:
  642. """
  643. Delete a string from the AutoCompleter index.
  644. Returns 1 if the string was found and deleted, 0 otherwise.
  645. For more information see `FT.SUGDEL <https://redis.io/commands/ft.sugdel>`_.
  646. """ # noqa
  647. return self.execute_command(SUGDEL_COMMAND, key, string)
  648. def sugget(
  649. self,
  650. key: str,
  651. prefix: str,
  652. fuzzy: bool = False,
  653. num: int = 10,
  654. with_scores: bool = False,
  655. with_payloads: bool = False,
  656. ) -> List[SuggestionParser]:
  657. """
  658. Get a list of suggestions from the AutoCompleter, for a given prefix.
  659. Parameters:
  660. prefix : str
  661. The prefix we are searching. **Must be valid ascii or utf-8**
  662. fuzzy : bool
  663. If set to true, the prefix search is done in fuzzy mode.
  664. **NOTE**: Running fuzzy searches on short (<3 letters) prefixes
  665. can be very
  666. slow, and even scan the entire index.
  667. with_scores : bool
  668. If set to true, we also return the (refactored) score of
  669. each suggestion.
  670. This is normally not needed, and is NOT the original score
  671. inserted into the index.
  672. with_payloads : bool
  673. Return suggestion payloads
  674. num : int
  675. The maximum number of results we return. Note that we might
  676. return less. The algorithm trims irrelevant suggestions.
  677. Returns:
  678. list:
  679. A list of Suggestion objects. If with_scores was False, the
  680. score of all suggestions is 1.
  681. For more information see `FT.SUGGET <https://redis.io/commands/ft.sugget>`_.
  682. """ # noqa
  683. args = [SUGGET_COMMAND, key, prefix, "MAX", num]
  684. if fuzzy:
  685. args.append(FUZZY)
  686. if with_scores:
  687. args.append(WITHSCORES)
  688. if with_payloads:
  689. args.append(WITHPAYLOADS)
  690. res = self.execute_command(*args)
  691. results = []
  692. if not res:
  693. return results
  694. parser = SuggestionParser(with_scores, with_payloads, res)
  695. return [s for s in parser]
  696. def synupdate(self, groupid: str, skipinitial: bool = False, *terms: List[str]):
  697. """
  698. Updates a synonym group.
  699. The command is used to create or update a synonym group with
  700. additional terms.
  701. Only documents which were indexed after the update will be affected.
  702. Parameters:
  703. groupid :
  704. Synonym group id.
  705. skipinitial : bool
  706. If set to true, we do not scan and index.
  707. terms :
  708. The terms.
  709. For more information see `FT.SYNUPDATE <https://redis.io/commands/ft.synupdate>`_.
  710. """ # noqa
  711. cmd = [SYNUPDATE_CMD, self.index_name, groupid]
  712. if skipinitial:
  713. cmd.extend(["SKIPINITIALSCAN"])
  714. cmd.extend(terms)
  715. return self.execute_command(*cmd)
  716. def syndump(self):
  717. """
  718. Dumps the contents of a synonym group.
  719. The command is used to dump the synonyms data structure.
  720. Returns a list of synonym terms and their synonym group ids.
  721. For more information see `FT.SYNDUMP <https://redis.io/commands/ft.syndump>`_.
  722. """ # noqa
  723. res = self.execute_command(SYNDUMP_CMD, self.index_name)
  724. return self._parse_results(SYNDUMP_CMD, res)
  725. class AsyncSearchCommands(SearchCommands):
  726. async def info(self):
  727. """
  728. Get info an stats about the the current index, including the number of
  729. documents, memory consumption, etc
  730. For more information see `FT.INFO <https://redis.io/commands/ft.info>`_.
  731. """
  732. res = await self.execute_command(INFO_CMD, self.index_name)
  733. return self._parse_results(INFO_CMD, res)
  734. async def search(
  735. self,
  736. query: Union[str, Query],
  737. query_params: Dict[str, Union[str, int, float]] = None,
  738. ):
  739. """
  740. Search the index for a given query, and return a result of documents
  741. ### Parameters
  742. - **query**: the search query. Either a text for simple queries with
  743. default parameters, or a Query object for complex queries.
  744. See RediSearch's documentation on query format
  745. For more information see `FT.SEARCH <https://redis.io/commands/ft.search>`_.
  746. """ # noqa
  747. args, query = self._mk_query_args(query, query_params=query_params)
  748. st = time.time()
  749. res = await self.execute_command(SEARCH_CMD, *args)
  750. if isinstance(res, Pipeline):
  751. return res
  752. return self._parse_results(
  753. SEARCH_CMD, res, query=query, duration=(time.time() - st) * 1000.0
  754. )
  755. async def aggregate(
  756. self,
  757. query: Union[str, Query],
  758. query_params: Dict[str, Union[str, int, float]] = None,
  759. ):
  760. """
  761. Issue an aggregation query.
  762. ### Parameters
  763. **query**: This can be either an `AggregateRequest`, or a `Cursor`
  764. An `AggregateResult` object is returned. You can access the rows from
  765. its `rows` property, which will always yield the rows of the result.
  766. For more information see `FT.AGGREGATE <https://redis.io/commands/ft.aggregate>`_.
  767. """ # noqa
  768. if isinstance(query, AggregateRequest):
  769. has_cursor = bool(query._cursor)
  770. cmd = [AGGREGATE_CMD, self.index_name] + query.build_args()
  771. elif isinstance(query, Cursor):
  772. has_cursor = True
  773. cmd = [CURSOR_CMD, "READ", self.index_name] + query.build_args()
  774. else:
  775. raise ValueError("Bad query", query)
  776. cmd += self.get_params_args(query_params)
  777. raw = await self.execute_command(*cmd)
  778. return self._parse_results(
  779. AGGREGATE_CMD, raw, query=query, has_cursor=has_cursor
  780. )
  781. async def spellcheck(self, query, distance=None, include=None, exclude=None):
  782. """
  783. Issue a spellcheck query
  784. ### Parameters
  785. **query**: search query.
  786. **distance***: the maximal Levenshtein distance for spelling
  787. suggestions (default: 1, max: 4).
  788. **include**: specifies an inclusion custom dictionary.
  789. **exclude**: specifies an exclusion custom dictionary.
  790. For more information see `FT.SPELLCHECK <https://redis.io/commands/ft.spellcheck>`_.
  791. """ # noqa
  792. cmd = [SPELLCHECK_CMD, self.index_name, query]
  793. if distance:
  794. cmd.extend(["DISTANCE", distance])
  795. if include:
  796. cmd.extend(["TERMS", "INCLUDE", include])
  797. if exclude:
  798. cmd.extend(["TERMS", "EXCLUDE", exclude])
  799. res = await self.execute_command(*cmd)
  800. return self._parse_results(SPELLCHECK_CMD, res)
  801. async def config_set(self, option: str, value: str) -> bool:
  802. """Set runtime configuration option.
  803. ### Parameters
  804. - **option**: the name of the configuration option.
  805. - **value**: a value for the configuration option.
  806. For more information see `FT.CONFIG SET <https://redis.io/commands/ft.config-set>`_.
  807. """ # noqa
  808. cmd = [CONFIG_CMD, "SET", option, value]
  809. raw = await self.execute_command(*cmd)
  810. return raw == "OK"
  811. async def config_get(self, option: str) -> str:
  812. """Get runtime configuration option value.
  813. ### Parameters
  814. - **option**: the name of the configuration option.
  815. For more information see `FT.CONFIG GET <https://redis.io/commands/ft.config-get>`_.
  816. """ # noqa
  817. cmd = [CONFIG_CMD, "GET", option]
  818. res = {}
  819. res = await self.execute_command(*cmd)
  820. return self._parse_results(CONFIG_CMD, res)
  821. async def load_document(self, id):
  822. """
  823. Load a single document by id
  824. """
  825. fields = await self.client.hgetall(id)
  826. f2 = {to_string(k): to_string(v) for k, v in fields.items()}
  827. fields = f2
  828. try:
  829. del fields["id"]
  830. except KeyError:
  831. pass
  832. return Document(id=id, **fields)
  833. async def sugadd(self, key, *suggestions, **kwargs):
  834. """
  835. Add suggestion terms to the AutoCompleter engine. Each suggestion has
  836. a score and string.
  837. If kwargs["increment"] is true and the terms are already in the
  838. server's dictionary, we increment their scores.
  839. For more information see `FT.SUGADD <https://redis.io/commands/ft.sugadd>`_.
  840. """ # noqa
  841. # If Transaction is not False it will MULTI/EXEC which will error
  842. pipe = self.pipeline(transaction=False)
  843. for sug in suggestions:
  844. args = [SUGADD_COMMAND, key, sug.string, sug.score]
  845. if kwargs.get("increment"):
  846. args.append("INCR")
  847. if sug.payload:
  848. args.append("PAYLOAD")
  849. args.append(sug.payload)
  850. pipe.execute_command(*args)
  851. return (await pipe.execute())[-1]
  852. async def sugget(
  853. self,
  854. key: str,
  855. prefix: str,
  856. fuzzy: bool = False,
  857. num: int = 10,
  858. with_scores: bool = False,
  859. with_payloads: bool = False,
  860. ) -> List[SuggestionParser]:
  861. """
  862. Get a list of suggestions from the AutoCompleter, for a given prefix.
  863. Parameters:
  864. prefix : str
  865. The prefix we are searching. **Must be valid ascii or utf-8**
  866. fuzzy : bool
  867. If set to true, the prefix search is done in fuzzy mode.
  868. **NOTE**: Running fuzzy searches on short (<3 letters) prefixes
  869. can be very
  870. slow, and even scan the entire index.
  871. with_scores : bool
  872. If set to true, we also return the (refactored) score of
  873. each suggestion.
  874. This is normally not needed, and is NOT the original score
  875. inserted into the index.
  876. with_payloads : bool
  877. Return suggestion payloads
  878. num : int
  879. The maximum number of results we return. Note that we might
  880. return less. The algorithm trims irrelevant suggestions.
  881. Returns:
  882. list:
  883. A list of Suggestion objects. If with_scores was False, the
  884. score of all suggestions is 1.
  885. For more information see `FT.SUGGET <https://redis.io/commands/ft.sugget>`_.
  886. """ # noqa
  887. args = [SUGGET_COMMAND, key, prefix, "MAX", num]
  888. if fuzzy:
  889. args.append(FUZZY)
  890. if with_scores:
  891. args.append(WITHSCORES)
  892. if with_payloads:
  893. args.append(WITHPAYLOADS)
  894. ret = await self.execute_command(*args)
  895. results = []
  896. if not ret:
  897. return results
  898. parser = SuggestionParser(with_scores, with_payloads, ret)
  899. return [s for s in parser]