Funktionierender Prototyp des Serious Games zur Vermittlung von Wissen zu Software-Engineering-Arbeitsmodellen.
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

search.py 11KB

1 year ago
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370
  1. import psycopg2
  2. from django.db.models import (
  3. CharField,
  4. Expression,
  5. Field,
  6. FloatField,
  7. Func,
  8. Lookup,
  9. TextField,
  10. Value,
  11. )
  12. from django.db.models.expressions import CombinedExpression, register_combinable_fields
  13. from django.db.models.functions import Cast, Coalesce
  14. class SearchVectorExact(Lookup):
  15. lookup_name = "exact"
  16. def process_rhs(self, qn, connection):
  17. if not isinstance(self.rhs, (SearchQuery, CombinedSearchQuery)):
  18. config = getattr(self.lhs, "config", None)
  19. self.rhs = SearchQuery(self.rhs, config=config)
  20. rhs, rhs_params = super().process_rhs(qn, connection)
  21. return rhs, rhs_params
  22. def as_sql(self, qn, connection):
  23. lhs, lhs_params = self.process_lhs(qn, connection)
  24. rhs, rhs_params = self.process_rhs(qn, connection)
  25. params = lhs_params + rhs_params
  26. return "%s @@ %s" % (lhs, rhs), params
  27. class SearchVectorField(Field):
  28. def db_type(self, connection):
  29. return "tsvector"
  30. class SearchQueryField(Field):
  31. def db_type(self, connection):
  32. return "tsquery"
  33. class SearchConfig(Expression):
  34. def __init__(self, config):
  35. super().__init__()
  36. if not hasattr(config, "resolve_expression"):
  37. config = Value(config)
  38. self.config = config
  39. @classmethod
  40. def from_parameter(cls, config):
  41. if config is None or isinstance(config, cls):
  42. return config
  43. return cls(config)
  44. def get_source_expressions(self):
  45. return [self.config]
  46. def set_source_expressions(self, exprs):
  47. (self.config,) = exprs
  48. def as_sql(self, compiler, connection):
  49. sql, params = compiler.compile(self.config)
  50. return "%s::regconfig" % sql, params
  51. class SearchVectorCombinable:
  52. ADD = "||"
  53. def _combine(self, other, connector, reversed):
  54. if not isinstance(other, SearchVectorCombinable):
  55. raise TypeError(
  56. "SearchVector can only be combined with other SearchVector "
  57. "instances, got %s." % type(other).__name__
  58. )
  59. if reversed:
  60. return CombinedSearchVector(other, connector, self, self.config)
  61. return CombinedSearchVector(self, connector, other, self.config)
  62. register_combinable_fields(
  63. SearchVectorField, SearchVectorCombinable.ADD, SearchVectorField, SearchVectorField
  64. )
  65. class SearchVector(SearchVectorCombinable, Func):
  66. function = "to_tsvector"
  67. arg_joiner = " || ' ' || "
  68. output_field = SearchVectorField()
  69. def __init__(self, *expressions, config=None, weight=None):
  70. super().__init__(*expressions)
  71. self.config = SearchConfig.from_parameter(config)
  72. if weight is not None and not hasattr(weight, "resolve_expression"):
  73. weight = Value(weight)
  74. self.weight = weight
  75. def resolve_expression(
  76. self, query=None, allow_joins=True, reuse=None, summarize=False, for_save=False
  77. ):
  78. resolved = super().resolve_expression(
  79. query, allow_joins, reuse, summarize, for_save
  80. )
  81. if self.config:
  82. resolved.config = self.config.resolve_expression(
  83. query, allow_joins, reuse, summarize, for_save
  84. )
  85. return resolved
  86. def as_sql(self, compiler, connection, function=None, template=None):
  87. clone = self.copy()
  88. clone.set_source_expressions(
  89. [
  90. Coalesce(
  91. expression
  92. if isinstance(expression.output_field, (CharField, TextField))
  93. else Cast(expression, TextField()),
  94. Value(""),
  95. )
  96. for expression in clone.get_source_expressions()
  97. ]
  98. )
  99. config_sql = None
  100. config_params = []
  101. if template is None:
  102. if clone.config:
  103. config_sql, config_params = compiler.compile(clone.config)
  104. template = "%(function)s(%(config)s, %(expressions)s)"
  105. else:
  106. template = clone.template
  107. sql, params = super(SearchVector, clone).as_sql(
  108. compiler,
  109. connection,
  110. function=function,
  111. template=template,
  112. config=config_sql,
  113. )
  114. extra_params = []
  115. if clone.weight:
  116. weight_sql, extra_params = compiler.compile(clone.weight)
  117. sql = "setweight({}, {})".format(sql, weight_sql)
  118. return sql, config_params + params + extra_params
  119. class CombinedSearchVector(SearchVectorCombinable, CombinedExpression):
  120. def __init__(self, lhs, connector, rhs, config, output_field=None):
  121. self.config = config
  122. super().__init__(lhs, connector, rhs, output_field)
  123. class SearchQueryCombinable:
  124. BITAND = "&&"
  125. BITOR = "||"
  126. def _combine(self, other, connector, reversed):
  127. if not isinstance(other, SearchQueryCombinable):
  128. raise TypeError(
  129. "SearchQuery can only be combined with other SearchQuery "
  130. "instances, got %s." % type(other).__name__
  131. )
  132. if reversed:
  133. return CombinedSearchQuery(other, connector, self, self.config)
  134. return CombinedSearchQuery(self, connector, other, self.config)
  135. # On Combinable, these are not implemented to reduce confusion with Q. In
  136. # this case we are actually (ab)using them to do logical combination so
  137. # it's consistent with other usage in Django.
  138. def __or__(self, other):
  139. return self._combine(other, self.BITOR, False)
  140. def __ror__(self, other):
  141. return self._combine(other, self.BITOR, True)
  142. def __and__(self, other):
  143. return self._combine(other, self.BITAND, False)
  144. def __rand__(self, other):
  145. return self._combine(other, self.BITAND, True)
  146. class SearchQuery(SearchQueryCombinable, Func):
  147. output_field = SearchQueryField()
  148. SEARCH_TYPES = {
  149. "plain": "plainto_tsquery",
  150. "phrase": "phraseto_tsquery",
  151. "raw": "to_tsquery",
  152. "websearch": "websearch_to_tsquery",
  153. }
  154. def __init__(
  155. self,
  156. value,
  157. output_field=None,
  158. *,
  159. config=None,
  160. invert=False,
  161. search_type="plain",
  162. ):
  163. self.function = self.SEARCH_TYPES.get(search_type)
  164. if self.function is None:
  165. raise ValueError("Unknown search_type argument '%s'." % search_type)
  166. if not hasattr(value, "resolve_expression"):
  167. value = Value(value)
  168. expressions = (value,)
  169. self.config = SearchConfig.from_parameter(config)
  170. if self.config is not None:
  171. expressions = (self.config,) + expressions
  172. self.invert = invert
  173. super().__init__(*expressions, output_field=output_field)
  174. def as_sql(self, compiler, connection, function=None, template=None):
  175. sql, params = super().as_sql(compiler, connection, function, template)
  176. if self.invert:
  177. sql = "!!(%s)" % sql
  178. return sql, params
  179. def __invert__(self):
  180. clone = self.copy()
  181. clone.invert = not self.invert
  182. return clone
  183. def __str__(self):
  184. result = super().__str__()
  185. return ("~%s" % result) if self.invert else result
  186. class CombinedSearchQuery(SearchQueryCombinable, CombinedExpression):
  187. def __init__(self, lhs, connector, rhs, config, output_field=None):
  188. self.config = config
  189. super().__init__(lhs, connector, rhs, output_field)
  190. def __str__(self):
  191. return "(%s)" % super().__str__()
  192. class SearchRank(Func):
  193. function = "ts_rank"
  194. output_field = FloatField()
  195. def __init__(
  196. self,
  197. vector,
  198. query,
  199. weights=None,
  200. normalization=None,
  201. cover_density=False,
  202. ):
  203. if not hasattr(vector, "resolve_expression"):
  204. vector = SearchVector(vector)
  205. if not hasattr(query, "resolve_expression"):
  206. query = SearchQuery(query)
  207. expressions = (vector, query)
  208. if weights is not None:
  209. if not hasattr(weights, "resolve_expression"):
  210. weights = Value(weights)
  211. expressions = (weights,) + expressions
  212. if normalization is not None:
  213. if not hasattr(normalization, "resolve_expression"):
  214. normalization = Value(normalization)
  215. expressions += (normalization,)
  216. if cover_density:
  217. self.function = "ts_rank_cd"
  218. super().__init__(*expressions)
  219. class SearchHeadline(Func):
  220. function = "ts_headline"
  221. template = "%(function)s(%(expressions)s%(options)s)"
  222. output_field = TextField()
  223. def __init__(
  224. self,
  225. expression,
  226. query,
  227. *,
  228. config=None,
  229. start_sel=None,
  230. stop_sel=None,
  231. max_words=None,
  232. min_words=None,
  233. short_word=None,
  234. highlight_all=None,
  235. max_fragments=None,
  236. fragment_delimiter=None,
  237. ):
  238. if not hasattr(query, "resolve_expression"):
  239. query = SearchQuery(query)
  240. options = {
  241. "StartSel": start_sel,
  242. "StopSel": stop_sel,
  243. "MaxWords": max_words,
  244. "MinWords": min_words,
  245. "ShortWord": short_word,
  246. "HighlightAll": highlight_all,
  247. "MaxFragments": max_fragments,
  248. "FragmentDelimiter": fragment_delimiter,
  249. }
  250. self.options = {
  251. option: value for option, value in options.items() if value is not None
  252. }
  253. expressions = (expression, query)
  254. if config is not None:
  255. config = SearchConfig.from_parameter(config)
  256. expressions = (config,) + expressions
  257. super().__init__(*expressions)
  258. def as_sql(self, compiler, connection, function=None, template=None):
  259. options_sql = ""
  260. options_params = []
  261. if self.options:
  262. # getquoted() returns a quoted bytestring of the adapted value.
  263. options_params.append(
  264. ", ".join(
  265. "%s=%s"
  266. % (
  267. option,
  268. psycopg2.extensions.adapt(value).getquoted().decode(),
  269. )
  270. for option, value in self.options.items()
  271. )
  272. )
  273. options_sql = ", %s"
  274. sql, params = super().as_sql(
  275. compiler,
  276. connection,
  277. function=function,
  278. template=template,
  279. options=options_sql,
  280. )
  281. return sql, params + options_params
  282. SearchVectorField.register_lookup(SearchVectorExact)
  283. class TrigramBase(Func):
  284. output_field = FloatField()
  285. def __init__(self, expression, string, **extra):
  286. if not hasattr(string, "resolve_expression"):
  287. string = Value(string)
  288. super().__init__(expression, string, **extra)
  289. class TrigramWordBase(Func):
  290. output_field = FloatField()
  291. def __init__(self, string, expression, **extra):
  292. if not hasattr(string, "resolve_expression"):
  293. string = Value(string)
  294. super().__init__(string, expression, **extra)
  295. class TrigramSimilarity(TrigramBase):
  296. function = "SIMILARITY"
  297. class TrigramDistance(TrigramBase):
  298. function = ""
  299. arg_joiner = " <-> "
  300. class TrigramWordDistance(TrigramWordBase):
  301. function = ""
  302. arg_joiner = " <<-> "
  303. class TrigramWordSimilarity(TrigramWordBase):
  304. function = "WORD_SIMILARITY"