Coverage for app / translator.py: 100%

1import re 

2from typing import Callable 

3 

4 

5def parse_query_delimiters(sql: str) -> list[str | None]: 

6 delimiters = re.findall(r"\s*;\s*", sql) 

7 return delimiters 

8 

9 

10def merge_sql_strings(queries: list[str], delimiters: list[str | None]) -> str: 

11 if not queries: 

12 return "" 

13 

14 if not delimiters: 

15 return queries[0] 

16 

17 n_queries = len(queries) 

18 n_delimiters = len(delimiters) 

19 diff = n_queries - n_delimiters 

20 

21 if n_delimiters != n_queries and n_delimiters != (n_queries - 1): 

22 raise ValueError( 

23 f"The number of delimiters ({n_delimiters}) must be equal to or one less than the " 

24 f"number of queries ({n_queries})." 

25 ) 

26 

27 result = "".join( 

28 [f"{query}{delimiter}" for query, delimiter in zip(queries, delimiters)] 

29 ) 

30 

31 if diff == 1: 

32 result += queries[-1] 

33 

34 return result 

35 

36 

37def get_casing(words: list[str]) -> Callable[[str], str]: 

38 first_alpha = next((char for char in "".join(words) if char.isalpha()), "a") 

39 return str.upper if first_alpha.isupper() else str.lower 

40 

41 

42def restore_casing(original: str, transpiled: str) -> str: 

43 """ 

44 SQLGlot doesn't preserve the original casing of the original sql string. 

45 This function does a 'best effort' job of trying to restore the orignal casing. 

46 

47 This is done by comparing each 'word' from the original sql string to the 

48 transpiled sql string. It looks for exact matches based on a case insensitive 

49 comparison. 

50 

51 We then figure out "implied" matches based on any words it skips. 

52 

53 Once the range of original words maps to the range of transpiled words 

54 we apply the casing of the original word based on the first letter. 

55 """ 

56 original_words = re.split(r"(\s+)", original) 

57 transpiled_words_lower = re.split(r"(\s+)", transpiled.lower()) 

58 original_mapping = [] 

59 transpiled_mapping = [] 

60 last_original_match_idx = 0 

61 last_transpiled_match_idx = 0 

62 for original_idx, original_word in enumerate(original_words): 

63 if original_word.isspace(): 

64 continue 

65 

66 original_word_lower = original_word.lower() 

67 transpiled_words_subset = transpiled_words_lower[last_transpiled_match_idx:] 

68 

69 if original_word_lower in transpiled_words_subset: 

70 transpiled_index = transpiled_words_subset.index(original_word_lower) 

71 transpiled_index += last_transpiled_match_idx 

72 

73 if (original_idx - last_original_match_idx) > 0: 

74 original_mapping.append( 

75 get_casing(original_words[last_original_match_idx:original_idx]) 

76 ) 

77 original_mapping.append( 

78 get_casing(original_words[original_idx : (original_idx + 1)]) 

79 ) 

80 last_original_match_idx = original_idx + 1 

81 

82 if (transpiled_index - last_transpiled_match_idx) > 0: 

83 transpiled_mapping.append((last_transpiled_match_idx, transpiled_index)) 

84 transpiled_mapping.append((transpiled_index, transpiled_index + 1)) 

85 last_transpiled_match_idx = transpiled_index + 1 

86 

87 if last_original_match_idx < len(original_words): 

88 original_mapping.append( 

89 get_casing(original_words[last_original_match_idx : len(original_words)]) 

90 ) 

91 

92 if last_transpiled_match_idx < len(transpiled_words_lower): 

93 transpiled_mapping.append( 

94 (last_transpiled_match_idx, len(transpiled_words_lower)) 

95 ) 

96 

97 result = "" 

98 for case_func, (start_idx, end_idx) in zip(original_mapping, transpiled_mapping): 

99 for word in transpiled_words_lower[start_idx:end_idx]: 

100 result += case_func(word) 

101 

102 return result