(parsed)
| 114 | |
| 115 | |
| 116 | def extract_column_names(parsed): |
| 117 | # Find the first DML token to check if it's a SELECT or INSERT/UPDATE/DELETE |
| 118 | idx, tok = parsed.token_next_by(t=DML) |
| 119 | tok_val = tok and tok.value.lower() |
| 120 | |
| 121 | if tok_val in ("insert", "update", "delete"): |
| 122 | # Jump ahead to the RETURNING clause where the list of column names is |
| 123 | idx, tok = parsed.token_next_by(idx, (Keyword, "returning")) |
| 124 | elif not tok_val == "select": |
| 125 | # Must be invalid CTE |
| 126 | return () |
| 127 | |
| 128 | # The next token should be either a column name, or a list of column names |
| 129 | idx, tok = parsed.token_next(idx, skip_ws=True, skip_cm=True) |
| 130 | return tuple(t.get_name() for t in _identifiers(tok)) |
| 131 | |
| 132 | |
| 133 | def token_start_pos(tokens, idx): |
no test coverage detected