rusackas commented on code in PR #42598:
URL: https://github.com/apache/superset/pull/42598#discussion_r3694002325
##########
tests/unit_tests/extensions/test_sqlalchemy.py:
##########
@@ -229,6 +230,108 @@ def test_superset_joins(
assert list(results) == [(10, "ten"), (20, "twenty")]
[email protected]
+def table1_large(session: Session, database1: "Database") -> Iterator[None]:
+ with database1.get_sqla_engine() as engine:
+ with engine.begin() as conn:
+ conn.execute(
+ text(
+ "CREATE TABLE table1_large (a INTEGER NOT NULL PRIMARY
KEY, "
+ "b INTEGER)"
+ )
+ )
+ conn.execute(
+ text("INSERT INTO table1_large (a, b) VALUES (1, 10), (2, 20),
(3, 30)")
+ )
+ db.session.commit()
+
+ yield
+
+ with engine.begin() as conn:
+ conn.execute(text("DROP TABLE table1_large"))
+ db.session.commit()
+
+
[email protected]
+def table2_late_match(session: Session, database2: "Database") ->
Iterator[None]:
+ with database2.get_sqla_engine() as engine:
+ with engine.begin() as conn:
+ conn.execute(
+ text(
+ "CREATE TABLE table2_late_match (a INTEGER NOT NULL
PRIMARY KEY, "
+ "b TEXT)"
+ )
+ )
+ conn.execute(
+ text("INSERT INTO table2_late_match (a, b) VALUES (3,
'thirty')")
+ )
+ db.session.commit()
+
+ yield
+
+ with engine.begin() as conn:
+ conn.execute(text("DROP TABLE table2_late_match"))
+ db.session.commit()
+
+
+@with_feature_flags(ENABLE_SUPERSET_META_DB=True)
+def test_superset_joins_with_limit_drops_matches(
+ mocker: MockerFixture,
+ monkeypatch: pytest.MonkeyPatch,
+ app_context: None,
+ table1_large: None,
+ table2_late_match: None,
+) -> None:
+ """
+ Regression for #36304: SUPERSET_META_DB_LIMIT is applied to each
+ underlying table independently, before the in-memory join runs. A row
+ that has a genuine match on the other side of the join but falls past
+ the per-table limit is silently dropped from the join result, with no
+ error or truncation warning.
+ """
+ # Use monkeypatch (rather than the `@with_config` decorator) so the
+ # config overrides are guaranteed to be undone even though this test is
+ # expected to fail its assertion until the underlying bug is fixed.
+ # `@with_config` only restores the original values after the wrapped
+ # test function returns normally, so an assertion failure here would
+ # otherwise leak SUPERSET_META_DB_LIMIT=2 into later tests.
+ monkeypatch.setitem(current_app.config, "DB_SQLA_URI_VALIDATOR", None)
+ monkeypatch.setitem(current_app.config, "SUPERSET_META_DB_LIMIT", 2)
+ monkeypatch.setitem(current_app.config, "DATABASE_OAUTH2_CLIENTS", {})
+ monkeypatch.setitem(current_app.config,
"SQLALCHEMY_CUSTOM_PASSWORD_STORE", None)
+
+ mocker.patch(
+ "superset.extensions.metadb.security_manager.raise_for_access",
+ return_value=None,
+ )
+
+ from flask import g
+
+ g.user = mocker.MagicMock()
+ g.user.is_anonymous = False
+
+ try:
+ engine = create_engine("superset://", future=True)
+ except Exception as e:
+ # Skip test if superset:// dialect can't be loaded (common in Docker)
+ pytest.skip(f"Superset dialect not available: {e}")
+
+ with engine.connect() as conn:
+ results = conn.execute(
+ text("""
+ SELECT t1.b, t2.b
+ FROM "database1.table1_large" AS t1
+ JOIN "database2.table2_late_match" AS t2
+ ON t1.a = t2.a
+ """)
+ )
+ # table2_late_match's only row (a=3) has a genuine match in
+ # table1_large (a=3, b=30), but SUPERSET_META_DB_LIMIT=2 truncates
+ # table1_large to its first two rows (a=1, a=2) before the join
+ # runs, so the join comes back empty instead of finding the match.
+ assert list(results) == [(30, "thirty")]
Review Comment:
Landed the adapter fix on top. SupersetAPSWDialect now tracks whether a
statement touches more than one table and skips the per-table cap when it does,
so the regression test should be green.
##########
superset/extensions/metadb.py:
##########
@@ -70,6 +73,18 @@
from superset import db, feature_flag_manager, security_manager
from superset.sql.parse import Table
+# Detects a `JOIN` keyword anywhere in the statement being executed against the
+# `superset://` engine. Shillelagh calls `SupersetShillelaghAdapter.get_data`
once
+# per underlying table, independently of any other table referenced by the same
+# statement, so it has no way on its own to tell whether it's being asked for a
+# standalone table or for one side of a join. `SupersetAPSWDialect.do_execute*`
+# populates `_executing_join_query` for the duration of a statement so that
+# `get_data` can tell the two cases apart (see `get_data` for why this
matters).
+_JOIN_KEYWORD_RE = re.compile(r"\bJOIN\b", re.IGNORECASE)
Review Comment:
Good catch, a comma join references two tables without the literal JOIN
keyword. Swapped the detection to count quoted db.table identifiers instead of
matching JOIN, covers both syntaxes now, and added a regression test for the
comma-join case.
##########
superset/extensions/metadb.py:
##########
@@ -119,6 +134,51 @@ def create_connect_args(self, url: URL) ->
tuple[tuple[()], dict[str, Any]]:
},
)
+ def do_execute(
+ self,
+ cursor: Any,
+ statement: str,
+ parameters: Any,
+ context: Any = None,
+ ) -> None:
+ with self._flag_join_query(statement):
+ super().do_execute(cursor, statement, parameters, context)
Review Comment:
Checked this empirically with a debug print inside get_data during the join
test, and it's called synchronously inside do_execute for this APSW/shillelagh
setup, the flag is still set when it runs. Not seeing this materialize in
practice, but flag it again if you spot a case where it does.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
