Extract constant table expresseions from a query Returns tuple (ctes, remainder_sql) ctes is a list of TableExpression namedtuples remainder_sql is the text from the original query after the CTEs have been stripped.
(sql)
| 47 | |
| 48 | |
| 49 | def extract_ctes(sql): |
| 50 | """ Extract constant table expresseions from a query |
| 51 | |
| 52 | Returns tuple (ctes, remainder_sql) |
| 53 | |
| 54 | ctes is a list of TableExpression namedtuples |
| 55 | remainder_sql is the text from the original query after the CTEs have |
| 56 | been stripped. |
| 57 | """ |
| 58 | |
| 59 | p = parse(sql)[0] |
| 60 | |
| 61 | # Make sure the first meaningful token is "WITH" which is necessary to |
| 62 | # define CTEs |
| 63 | idx, tok = p.token_next(-1, skip_ws=True, skip_cm=True) |
| 64 | if not (tok and tok.ttype == CTE): |
| 65 | return [], sql |
| 66 | |
| 67 | # Get the next (meaningful) token, which should be the first CTE |
| 68 | idx, tok = p.token_next(idx) |
| 69 | if not tok: |
| 70 | return ([], '') |
| 71 | start_pos = token_start_pos(p.tokens, idx) |
| 72 | ctes = [] |
| 73 | |
| 74 | if isinstance(tok, IdentifierList): |
| 75 | # Multiple ctes |
| 76 | for t in tok.get_identifiers(): |
| 77 | cte_start_offset = token_start_pos(tok.tokens, tok.token_index(t)) |
| 78 | cte = get_cte_from_token(t, start_pos + cte_start_offset) |
| 79 | if not cte: |
| 80 | continue |
| 81 | ctes.append(cte) |
| 82 | elif isinstance(tok, Identifier): |
| 83 | # A single CTE |
| 84 | cte = get_cte_from_token(tok, start_pos) |
| 85 | if cte: |
| 86 | ctes.append(cte) |
| 87 | |
| 88 | idx = p.token_index(tok) + 1 |
| 89 | |
| 90 | # Collapse everything after the ctes into a remainder query |
| 91 | remainder = u''.join(str(tok) for tok in p.tokens[idx:]) |
| 92 | |
| 93 | return ctes, remainder |
| 94 | |
| 95 | |
| 96 | def get_cte_from_token(tok, pos0): |