From 4251b4e29405b21fad32ddc70807a04696c69374 Mon Sep 17 00:00:00 2001
From: Ville Brofeldt <33317356+villebro@users.noreply.github.com>
Date: Wed, 26 Aug 2020 11:55:35 +0300
Subject: [PATCH] feat(row-level-security): add hook for customizing form
 dropdowns (#10683)

---
 superset/config.py                | 10 ++++++++++
 superset/connectors/sqla/views.py |  3 +++
 2 files changed, 13 insertions(+)

diff --git a/superset/config.py b/superset/config.py
index db5aa6e085..0dfde99054 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -869,6 +869,16 @@ TALISMAN_CONFIG = {
 # a custom security config could potentially give access to setting filters on
 # tables that users do not have access to.
 ENABLE_ROW_LEVEL_SECURITY = False
+# It is possible to customize which tables and roles are featured in the RLS
+# dropdown. When set, this dict is assigned to `add_form_query_rel_fields` and
+# `edit_form_query_rel_fields` on `RowLevelSecurityFiltersModelView`. Example:
+#
+# from flask_appbuilder.models.sqla import filters
+# RLS_FORM_QUERY_REL_FIELDS = {
+#     "roles": [["name", filters.FilterStartsWith, "RlsRole"]]
+#     "tables": [["table_name", filters.FilterContains, "rls"]]
+# }
+RLS_FORM_QUERY_REL_FIELDS: Optional[Dict[str, List[List[Any]]]] = None
 
 #
 # Flask session cookie options
diff --git a/superset/connectors/sqla/views.py b/superset/connectors/sqla/views.py
index 107e10abaf..23f4c1c793 100644
--- a/superset/connectors/sqla/views.py
+++ b/superset/connectors/sqla/views.py
@@ -274,6 +274,9 @@ class RowLevelSecurityFiltersModelView(  # pylint: disable=too-many-ancestors
         "creator": _("Creator"),
         "modified": _("Modified"),
     }
+    if app.config["RLS_FORM_QUERY_REL_FIELDS"]:
+        add_form_query_rel_fields = app.config["RLS_FORM_QUERY_REL_FIELDS"]
+        edit_form_query_rel_fields = add_form_query_rel_fields
 
 
 class TableModelView(  # pylint: disable=too-many-ancestors