def _collect_tables( cls, statement: token_groups.Statement) -> typing.List[sql_table.Table]: idx, _ = statement.token_next_by(m=[ (token_types.Keyword, "FROM"), (token_types.Keyword, "INTO"), (token_types.DML, "UPDATE"), ]) _, maybe_table_identifier = statement.token_next(idx=idx, skip_cm=True, skip_ws=True) if isinstance(maybe_table_identifier, token_groups.Function): maybe_table_identifier = maybe_table_identifier.token_first( skip_cm=True, skip_ws=True) # If we can't find a single table identifier, it means that multiple tables # are referenced in the FROM/INTO clause, which isn't supported. if not isinstance(maybe_table_identifier, token_groups.Identifier): raise exceptions.NotSupportedError( "In order to query multiple tables at a time, you must join them " "together with a JOIN clause.") table_identifier = maybe_table_identifier tables = [sql_table.Table.from_identifier(table_identifier)] while True: idx, join_kw = statement.token_next_by(m=(token_types.Keyword, "JOIN"), idx=idx) if join_kw is None: break idx, table_identifier = statement.token_next(idx, skip_ws=True, skip_cm=True) table = sql_table.Table.from_identifier(table_identifier) idx, comparison_group = statement.token_next_by( i=token_groups.Comparison, idx=idx) table.add_join(tables[-1], comparison_group, sql_table.JoinDirection.LEFT) tables.append(table) return tables
def get_token_next(statement: Statement, t: TokenType) -> TokenType: """`statement`中のあるトークン`t`の次のトークンを取得。コメントと空白はスキップ。""" if isinstance(t, ExtraToken): t = t.tokens[-1] return statement.token_next( statement.token_index(t), skip_ws=True, skip_cm=True )[1]
def from_statement( cls, statement: token_groups.Statement) -> typing.Optional[OrderBy]: """Extract results ordering from an SQL statement. Params: ------- statement: A full SQL statement Returns: -------- An OrderBy object with the SQL ORDER BY attributes. """ idx, order_by = statement.token_next_by(m=(token_types.Keyword, "ORDER BY")) if order_by is None: return None idx, identifier = statement.token_next(skip_cm=True, skip_ws=True, idx=idx) direction = cls._extract_direction(identifier) if direction is None: columns = sql_table.Column.from_identifier_group(identifier) else: # Because of how sqlparse erroneously groups the final column identifier # with the direction keyword, we have to parse identifiers separately, # drilling down an extra level for the final token. nested_columns = [ sql_table.Column.from_identifier_group(token) for token in identifier.tokens[:-1] if isinstance(token, (token_groups.Identifier, token_groups.IdentifierList)) ] # If we order by a single column, the final token will be the # direction keyword token. Otherwise, it will be an identifier with both # the final column identifier and the direction keyword. maybe_column_identifier = identifier.tokens[-1] if maybe_column_identifier.is_group: column_identifier = maybe_column_identifier _, final_column_identifier = column_identifier.token_next_by( i=token_groups.Identifier) nested_columns.append( sql_table.Column.from_identifier_group( final_column_identifier)) columns = list(itertools.chain.from_iterable(nested_columns)) return cls(columns=columns, direction=direction)
def _build_select_query(cls, statement: token_groups.Statement, tables: typing.List[sql_table.Table]) -> SQLQuery: _, wildcard = statement.token_next_by(t=(token_types.Wildcard)) if wildcard is not None: raise exceptions.NotSupportedError( "Wildcards ('*') are not yet supported") _, identifiers = statement.token_next_by(i=( token_groups.Identifier, token_groups.IdentifierList, token_groups.Function, )) for column in sql_table.Column.from_identifier_group(identifiers): try: table = next(table for table in tables if table.name == column.table_name) except StopIteration: table = tables[0] table.add_column(column) _, distinct = statement.token_next_by(m=(token_types.Keyword, "DISTINCT")) idx, _ = statement.token_next_by(m=(token_types.Keyword, "LIMIT")) _, limit = statement.token_next(skip_cm=True, skip_ws=True, idx=idx) limit_value = None if limit is None else int(limit.value) order_by = OrderBy.from_statement(statement) return cls( str(statement), tables=tables, distinct=bool(distinct), order_by=order_by, limit=limit_value, )