betodealmeida commented on a change in pull request #19454: URL: https://github.com/apache/superset/pull/19454#discussion_r839825075
########## File path: superset/sql_parse.py ########## @@ -539,45 +540,67 @@ def add_table_name(rls: TokenList, table: str) -> None: tokens.extend(token.tokens) -def matches_table_name(candidate: Token, table: str) -> bool: +def get_rls_for_table( + candidate: Token, + database_id: int, + default_schema: Optional[str], +) -> Optional[TokenList]: """ - Returns if the token represents a reference to the table. - - Tables can be fully qualified with periods. - - Note that in theory a table should be represented as an identifier, but due to - sqlparse's aggressive list of keywords (spanning multiple dialects) often it gets - classified as a keyword. + Given a table name, return any associated RLS predicates. """ + # pylint: disable=import-outside-toplevel + from superset import db + from superset.connectors.sqla.models import SqlaTable + if not isinstance(candidate, Identifier): candidate = Identifier([Token(Name, candidate.value)]) - target = sqlparse.parse(table)[0].tokens[0] - if not isinstance(target, Identifier): - target = Identifier([Token(Name, target.value)]) + table = ParsedQuery._get_table(candidate) # pylint: disable=protected-access + if not table: + return None - # match from right to left, splitting on the period, eg, schema.table == table - for left, right in zip(candidate.tokens[::-1], target.tokens[::-1]): - if left.value != right.value: - return False + dataset = ( + db.session.query(SqlaTable) + .filter( + and_( + SqlaTable.database_id == database_id, + SqlaTable.schema == (table.schema or default_schema), + SqlaTable.table_name == table.table, + ) + ) + .one_or_none() + ) + if not dataset: + return None - return True + template_processor = dataset.get_template_processor() + # pylint: disable=protected-access + predicate = " AND ".join( + str(filter_) + for filter_ in dataset._get_sqla_row_level_filters(template_processor) + ) + rls = sqlparse.parse(predicate)[0] + add_table_name(rls, str(dataset)) + return rls Review comment: Ah, good point! Fixed it. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: notifications-unsubscr...@superset.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: notifications-unsubscr...@superset.apache.org For additional commands, e-mail: notifications-h...@superset.apache.org