from marshmallow import fields from analytics.connectors.snowflake import AbstractSnowflakeQuery from analytics.queries.schema import QueryWithPermissionsSchema class ExampleSchema(QueryWithPermissionsSchema): test = fields.Str() table = fields.Str() field = fields.Str() number = fields.Int() class ExampleSnowflakeQueryUsingPermissionsMacro(AbstractSnowflakeQuery): filename = "test_template_using_permissions_macro.sql" query_schema = ExampleSchema class ExampleSnowflakeQueryUsingPermissionsMacroV2(AbstractSnowflakeQuery): filename = "test_template_using_permissions_macro_v2.sql" query_schema = ExampleSchema class ExampleSnowflakeQueryUsingPermissionsMacroRollupGrain(AbstractSnowflakeQuery): filename = "test_template_using_permissions_macro_rollup_grain.sql" query_schema = ExampleSchema class ExampleSnowflakeQueryUsingPermissionsMacroScopedRollupGrain( AbstractSnowflakeQuery ): filename = "test_template_using_permissions_macro_scoped_rollup_grain.sql" query_schema = ExampleSchema def strip_sql(sql): if hasattr(sql, "text"): sql = sql.text return " ".join(sql.split()) def permission_test(mock_execute_orm, permissions, expected_call_args): q = ExampleSnowflakeQueryUsingPermissionsMacro({"field": "test", **permissions}) q.execute() mock_execute_orm.reset() sql, args = mock_execute_orm.call_args[0] expected_sql, expected_args = expected_call_args sql = strip_sql(sql) assert sql == expected_sql assert args == expected_args def test_full_access_permissions(mock_execute_orm): permission_test( mock_execute_orm, { "permission_label_ids": [], "permission_artist_ids": [], "permission_subaccount_ids": [], "permission_label_participant_ids": [], }, ( "SELECT :field_1 FROM test_table where true ;", { "field_1": "test", }, ), ) def test_no_access_permissions(mock_execute_orm): permission_test( mock_execute_orm, { "permission_label_ids": None, "permission_artist_ids": None, "permission_subaccount_ids": None, "permission_label_participant_ids": None, }, ( "SELECT :field_1 FROM test_table where false ;", { "field_1": "test", }, ), ) def test_artist_id_access_permissions(mock_execute_orm): permission_test( mock_execute_orm, { "permission_label_ids": None, "permission_artist_ids": [1, 2, 3], "permission_subaccount_ids": None, "permission_label_participant_ids": None, }, ( "SELECT :field_1 FROM test_table " "where s.product_id IN " "( SELECT product_id FROM dim_release " "WHERE artistid IN (:inclause_2,:inclause_3,:inclause_4) ) ;", {"field_1": "test", "inclause_2": 1, "inclause_3": 2, "inclause_4": 3}, ), ) def test_label_artist_id_access_permissions(mock_execute_orm): permission_test( mock_execute_orm, { "permission_label_ids": [1, 2, 3], "permission_artist_ids": [1, 2, 3], "permission_subaccount_ids": None, "permission_label_participant_ids": None, }, ( "SELECT :field_1 FROM test_table where s.product_id IN " "( SELECT product_id FROM dim_release " "WHERE labelid IN (:inclause_2,:inclause_3,:inclause_4) " "OR artistid IN (:inclause_5,:inclause_6,:inclause_7) ) ;", { "field_1": "test", "inclause_2": 1, "inclause_3": 2, "inclause_4": 3, "inclause_5": 1, "inclause_6": 2, "inclause_7": 3, }, ), ) def test_subaccount_artist_id_access_permissions(mock_execute_orm): permission_test( mock_execute_orm, { "permission_label_ids": None, "permission_artist_ids": [1, 2, 3], "permission_subaccount_ids": [4, 5, 6], "permission_label_participant_ids": None, }, ( "SELECT :field_1 FROM test_table where s.product_id IN" " ( SELECT product_id FROM dim_release " "WHERE subaccountid IN (:inclause_2,:inclause_3,:inclause_4) " "OR artistid IN (:inclause_5,:inclause_6,:inclause_7) ) ;", { "field_1": "test", "inclause_2": 4, "inclause_3": 5, "inclause_4": 6, "inclause_5": 1, "inclause_6": 2, "inclause_7": 3, }, ), ) def test_label_participant_ids_access_permissions(mock_execute_orm): permission_test( mock_execute_orm, { "permission_label_ids": None, "permission_artist_ids": [1, 2, 3], "permission_subaccount_ids": None, "permission_label_participant_ids": [4, 5, 6], }, ( "SELECT :field_1 FROM test_table where s.product_id IN " "( SELECT product_id FROM dim_release WHERE" " artistid IN (:inclause_2,:inclause_3,:inclause_4) " "UNION SELECT product_id FROM " "label_participant_participated_in_orchard_product WHERE " "label_participant_id IN (:inclause_5,:inclause_6,:inclause_7) " "AND participated_as in" " ('performer', 'featuring', 'remixer') ) ;", { "field_1": "test", "inclause_2": 1, "inclause_3": 2, "inclause_4": 3, "inclause_5": 4, "inclause_6": 5, "inclause_7": 6, }, ), ) def test_full_access_with_feed_ids_permissions(mock_execute_orm): permission_test( mock_execute_orm, { "permission_label_ids": [], "permission_artist_ids": [], "permission_subaccount_ids": [], "permission_label_participant_ids": [], "permission_feed_ids": [1, 2, 3], }, ( "SELECT :field_1 FROM test_table where true AND feed_id IN (:inclause_2,:inclause_3,:inclause_4) ;", # noqa {"field_1": "test", "inclause_2": 1, "inclause_3": 2, "inclause_4": 3}, ), ) def _permission_test_v2(mock_execute_orm, permissions, expected_fragments): """Render the v2 macro template and assert each expected fragment appears.""" q = ExampleSnowflakeQueryUsingPermissionsMacroV2({"field": "test", **permissions}) q.execute() mock_execute_orm.reset() sql, _args = mock_execute_orm.call_args[0] sql = strip_sql(sql) for fragment in expected_fragments: assert ( fragment in sql ), f"expected fragment not found in rendered SQL:\n {fragment}\n--\nsql:\n {sql}" def test_label_participant_ids_access_permissions_v2(mock_execute_orm): """v2 LP branch: EXISTS join against product_ownership_access with the `access_until_date` gate. Regression guard against the post-transfer backdoor — the rendered SQL MUST time-slice the LP arm. The LP permission renders TWO arms: the current PARTICIPATED_IN arm and the historical USED_TO_PARTICIPATE_IN arm that re-grants a moved-away artist its former-owner slice (see test_label_participant_used_to_*).""" _permission_test_v2( mock_execute_orm, { "permission_label_ids": None, "permission_artist_ids": None, "permission_subaccount_ids": None, "permission_label_participant_ids": [4, 5, 6], }, [ "EXISTS (", "FROM label_participant_participated_in_orchard_product lp", "JOIN product_ownership_access poa", "ON poa.product_id = lp.product_id", "WHERE lp.product_id = s.product_id", "lp.participated_as IN ('performer', 'featuring', 'remixer')", "s.download_activity_date <= COALESCE(poa.access_until_date, DATE '9999-12-31')", # historical (USED_TO) arm coexists with the current arm "FROM label_participant_used_to_participate_in_orchard_product lp", "poa.access_until_date IS NOT NULL", ], ) def test_artist_ids_access_permissions_v2(mock_execute_orm): """v2 artist branch: EXISTS join against product_ownership_access with the `access_until_date` gate.""" _permission_test_v2( mock_execute_orm, { "permission_label_ids": None, "permission_artist_ids": [1, 2, 3], "permission_subaccount_ids": None, "permission_label_participant_ids": None, }, [ "EXISTS (", "FROM dim_release dr", "JOIN product_ownership_access poa", "ON poa.product_id = dr.product_id", "WHERE dr.product_id = s.product_id", "dr.artistid IN", "s.download_activity_date <= COALESCE(poa.access_until_date, DATE '9999-12-31')", ], ) def test_all_arms_combined_v2(mock_execute_orm): """v2 with labels + artist + LP: FOUR separate EXISTS clauses, all time-sliced, OR-joined at the top (no UNION of arms anymore). The LP permission contributes two arms — the current PARTICIPATED_IN arm and the historical USED_TO_PARTICIPATE_IN arm — so labels + artist + LP yields 1 + 1 + 2 = 4 EXISTS clauses.""" q = ExampleSnowflakeQueryUsingPermissionsMacroV2( { "field": "test", "permission_label_ids": [1], "permission_artist_ids": [2], "permission_subaccount_ids": None, "permission_label_participant_ids": [3], } ) q.execute() sql = strip_sql(mock_execute_orm.call_args[0][0]) assert sql.count("EXISTS (") == 4 # every arm must carry the access_until_date gate assert sql.count("poa.access_until_date") >= 4 # artist + LP arms no longer fold into a single IN-subquery via UNION assert "UNION" not in sql # the LP permission contributes both the current and the historical table assert "FROM label_participant_participated_in_orchard_product lp" in sql assert "FROM label_participant_used_to_participate_in_orchard_product lp" in sql def test_full_access_rollup_grain_macro(mock_execute_orm): """Employees on a widened `_ROLLUP` table: macro must emit an EXISTS-on-`is_current = TRUE` correlated by (product_id, label_id, subaccount_id) instead of `true`. Without this guard the rollup's one-row-per-ownership-tuple grain would double-count the pre-transfer window for any product that has ever been transferred. """ q = ExampleSnowflakeQueryUsingPermissionsMacroRollupGrain( { "field": "test", "permission_label_ids": [], "permission_artist_ids": [], "permission_subaccount_ids": [], "permission_label_participant_ids": [], } ) q.execute() sql = strip_sql(mock_execute_orm.call_args[0][0]) assert "EXISTS (" in sql assert "FROM product_ownership_access poa" in sql assert "poa.product_id = s.product_id" in sql assert "poa.label_id IS NOT DISTINCT FROM s.label_id" in sql assert "poa.subaccount_id IS NOT DISTINCT FROM s.subaccount_id" in sql assert "poa.is_current = TRUE" in sql # No standalone `where true` short-circuit on widened rollups. assert "where true" not in sql.lower() def test_scoped_label_rollup_grain_macro(mock_execute_orm): """Scoped users (label_ids) on a widened `_ROLLUP` table: legacy `product_id IN (dim_release …)` matches every ownership-tuple row of the user's products, so for any transferred UPC the SUM picks up both the current-owner row and the former-owner row, double-counting the pre-cutoff window. The macro must emit an EXISTS that pins `s.(label_id, subaccount_id)` to a poa row that is also accessible to the user.""" q = ExampleSnowflakeQueryUsingPermissionsMacroScopedRollupGrain( { "field": "test", "permission_label_ids": [1, 2, 3], "permission_artist_ids": None, "permission_subaccount_ids": None, "permission_label_participant_ids": None, } ) q.execute() sql = strip_sql(mock_execute_orm.call_args[0][0]) assert "EXISTS (" in sql assert "FROM product_ownership_access poa" in sql assert "poa.label_id IN (" in sql # The (product_id, label_id, subaccount_id) correlation to `s` is factored # to the outer EXISTS over the per-arm UNION ALL (acc), not per-arm poa. assert "acc.product_id = s.product_id" in sql assert "acc.label_id IS NOT DISTINCT FROM s.label_id" in sql assert "acc.subaccount_id IS NOT DISTINCT FROM s.subaccount_id" in sql # No tautology `IN (dim_release …)` shape — pre-fix bug. assert "dim_release" not in sql.lower() # Time-slice predicate falls away on rollups (no download_activity_date). assert "access_until_date" not in sql.lower() assert "download_activity_date" not in sql.lower() def test_scoped_subaccount_rollup_grain_macro(mock_execute_orm): """subaccount_ids arm renders the same EXISTS shape, keyed on `poa.subaccount_id` instead of `poa.label_id`.""" q = ExampleSnowflakeQueryUsingPermissionsMacroScopedRollupGrain( { "field": "test", "permission_label_ids": None, "permission_artist_ids": None, "permission_subaccount_ids": [4, 5, 6], "permission_label_participant_ids": None, } ) q.execute() sql = strip_sql(mock_execute_orm.call_args[0][0]) assert "EXISTS (" in sql assert "poa.subaccount_id IN (" in sql assert "acc.label_id IS NOT DISTINCT FROM s.label_id" in sql assert "acc.subaccount_id IS NOT DISTINCT FROM s.subaccount_id" in sql def test_scoped_artist_rollup_grain_macro(mock_execute_orm): """Artist arm on a widened `_ROLLUP`: dim_release reflects only the *current* artist (no transfer history), so pin to the current-owner poa row via `poa.is_current = TRUE`. Same shape as the employee `_full_access_rollup_filter` but gated by the artist permission.""" q = ExampleSnowflakeQueryUsingPermissionsMacroScopedRollupGrain( { "field": "test", "permission_label_ids": None, "permission_artist_ids": [777], "permission_subaccount_ids": None, "permission_label_participant_ids": None, } ) q.execute() sql = strip_sql(mock_execute_orm.call_args[0][0]) assert "EXISTS (" in sql assert "FROM dim_release dr" in sql assert "JOIN product_ownership_access poa" in sql assert "ON poa.product_id = dr.product_id" in sql assert "dr.artistid IN (" in sql assert "poa.is_current = TRUE" in sql def test_scoped_label_participant_rollup_grain_macro(mock_execute_orm): """LP arm on a widened `_ROLLUP`: the current PARTICIPATED_IN arm keeps the is_current=TRUE pin (`lp` only knows the current product set), and a second historical USED_TO_PARTICIPATE_IN arm matches the former-owner window's rollup rows (access_until_date IS NOT NULL).""" q = ExampleSnowflakeQueryUsingPermissionsMacroScopedRollupGrain( { "field": "test", "permission_label_ids": None, "permission_artist_ids": None, "permission_subaccount_ids": None, "permission_label_participant_ids": [10000], } ) q.execute() sql = strip_sql(mock_execute_orm.call_args[0][0]) assert "EXISTS (" in sql assert "FROM label_participant_participated_in_orchard_product lp" in sql assert "JOIN product_ownership_access poa" in sql assert "ON poa.product_id = lp.product_id" in sql assert "lp.label_participant_id IN (" in sql assert "poa.is_current = TRUE" in sql # historical (USED_TO) rollup arm: former-owner window only, pinned to # the originating account — no forward leak. assert "FROM label_participant_used_to_participate_in_orchard_product lp" in sql assert "poa.label_id IS NOT DISTINCT FROM lp.vendor_id" in sql assert "poa.access_until_date IS NOT NULL" in sql # Perf regression guard (root cause: OR-of-EXISTS -> SYSTEM$DISTINCT_ID # re-aggregation -> 60s rollup timeout). The arms MUST collapse into ONE # EXISTS over a UNION ALL, never an `EXISTS(...) OR EXISTS(...)` chain. assert sql.count("EXISTS (") == 1 assert "UNION ALL" in sql assert "OR EXISTS" not in sql def test_scoped_full_access_v2_rollup_grain_macro(mock_execute_orm): """Full-access keeps its own current-only filter even on the v2 template — rollup_grain switches the macro shape regardless of `transfer_product_ownership_enabled`.""" q = ExampleSnowflakeQueryUsingPermissionsMacroScopedRollupGrain( { "field": "test", "permission_label_ids": [], "permission_artist_ids": [], "permission_subaccount_ids": [], "permission_label_participant_ids": [], } ) q.execute() sql = strip_sql(mock_execute_orm.call_args[0][0]) assert "poa.is_current = TRUE" in sql assert "poa.label_id IS NOT DISTINCT FROM s.label_id" in sql def test_label_participant_used_to_participate_in_v2(mock_execute_orm): """v2 DAILY historical arm: the USED_TO_PARTICIPATE_IN shadow re-grants a moved-away artist access to the transferred product, frozen to their former-owner window. The arm MUST: - read label_participant_used_to_participate_in_orchard_product, - scope poa to the originating account (vendor_id/subaccount_id), - admit only the closed former-owner window (access_until_date IS NOT NULL) and time-slice the day (download_activity_date <= cutoff) so there is no forward leak past the move boundary, - keep the same participated_as role filter as the current arm. """ _permission_test_v2( mock_execute_orm, { "permission_label_ids": None, "permission_artist_ids": None, "permission_subaccount_ids": None, "permission_label_participant_ids": [4, 5, 6], }, [ "FROM label_participant_used_to_participate_in_orchard_product lp", "JOIN product_ownership_access poa", "ON poa.product_id = lp.product_id", "poa.label_id IS NOT DISTINCT FROM lp.vendor_id", # 0 (USED_TO node key) and NULL (poa) both mean "no subaccount". "COALESCE(poa.subaccount_id, 0) = COALESCE(lp.subaccount_id, 0)", "lp.label_participant_id IN", "lp.participated_as IN ('performer', 'featuring', 'remixer')", "poa.access_until_date IS NOT NULL", "s.download_activity_date <= poa.access_until_date", ], ) def test_label_participant_used_to_participate_in_rollup_grain(mock_execute_orm): """v2 ROLLUP historical arm: on the widened rollup there is no download_activity_date, so the freeze is expressed by matching only the former-owner window's rollup rows — access_until_date IS NOT NULL, with the originating-account pin tying poa.label_id/subaccount_id to BOTH the USED_TO row's vendor and the rollup row's owner columns.""" q = ExampleSnowflakeQueryUsingPermissionsMacroScopedRollupGrain( { "field": "test", "permission_label_ids": None, "permission_artist_ids": None, "permission_subaccount_ids": None, "permission_label_participant_ids": [10000], } ) q.execute() sql = strip_sql(mock_execute_orm.call_args[0][0]) assert "FROM label_participant_used_to_participate_in_orchard_product lp" in sql assert "poa.label_id IS NOT DISTINCT FROM lp.vendor_id" in sql # 0 (USED_TO node key) and NULL (poa) both mean "no subaccount"; the # originating-account pin normalizes both before comparing. assert "COALESCE(poa.subaccount_id, 0) = COALESCE(lp.subaccount_id, 0)" in sql # the rollup-row owner pin (null-safe, both NULL) is factored to the outer # acc correlation over the per-arm UNION ALL. assert "acc.label_id IS NOT DISTINCT FROM s.label_id" in sql assert "acc.subaccount_id IS NOT DISTINCT FROM s.subaccount_id" in sql assert "poa.access_until_date IS NOT NULL" in sql # rollup has no per-day column — the historical arm must NOT invent one assert "download_activity_date" not in sql.lower()