Coverage for requests_tracker/sql/sql_parser.py: 100%

1from functools import lru_cache 

2from typing import TYPE_CHECKING, Generator, Tuple 

3 

4import sqlparse # type: ignore 

5from django.utils.html import escape 

6 

7if TYPE_CHECKING: 

8 SQLParseFilterGenerator = Generator[Tuple[sqlparse.tokens.Token, str], None, None] 

9 

10 

11class BoldKeywordFilter: 

12 """sqlparse filter to bold SQ = L keywords""" 

13 

14 @staticmethod 

15 def process(stream: "SQLParseFilterGenerator") -> "SQLParseFilterGenerator": 

16 """Process the token stream""" 

17 for token_type, value in stream: 

18 is_keyword = token_type in sqlparse.tokens.Keyword 

19 if is_keyword: 

20 yield sqlparse.tokens.Text, "<strong>" 

21 yield token_type, escape(value) 

22 if is_keyword: 

23 yield sqlparse.tokens.Text, "</strong>" 

24 

25 

26@lru_cache(maxsize=128) 

27def parse_sql(sql: str, align_indent: bool) -> str: 

28 stack = get_filter_stack(aligned_indent=align_indent) 

29 return "".join(stack.run(sql)) 

30 

31 

32@lru_cache(maxsize=None) 

33def get_filter_stack( 

34 aligned_indent: bool, 

35) -> sqlparse.engine.FilterStack: 

36 stack = sqlparse.engine.FilterStack() 

37 stack.enable_grouping() 

38 if aligned_indent: 

39 stack.stmtprocess.append( 

40 sqlparse.filters.AlignedIndentFilter(char="&nbsp;", n="<br/>") 

41 ) 

42 stack.preprocess.append(BoldKeywordFilter()) # add our custom filter 

43 stack.postprocess.append(sqlparse.filters.SerializerUnicode()) # tokens -> strings 

44 return stack