Support for excluding schemas

getredash · Jul 22, 2021 · dfaa6ff · dfaa6ff
1 parent 92e5d78
commit dfaa6ff
Showing 1 changed file with 21 additions and 3 deletions.
diff --git a/redash/query_runner/pg.py b/redash/query_runner/pg.py
@@ -341,6 +341,11 @@ def configuration_schema(cls):
                     "title": "Query Group for Scheduled Queries",
                     "default": "default",
                 },
+                "exclude_schemas": {
+                    "type": "string",
+                    "title": "Exclude schemas from scanning tables",
+                    "default": "pg_internal,pg_catalog,information_schema,pg_temp_%",
+                },
             },
             "order": [
                 "host",
@@ -349,6 +354,7 @@ def configuration_schema(cls):
                 "password",
                 "dbname",
                 "sslmode",
+                "exclude_schemas",
                 "adhoc_query_group",
                 "scheduled_query_group",
             ],
@@ -378,15 +384,18 @@ def _get_tables(self, schema):
         # https://docs.aws.amazon.com/redshift/latest/dg/r_HAS_SCHEMA_PRIVILEGE.html
         # https://docs.aws.amazon.com/redshift/latest/dg/r_SVV_EXTERNAL_SCHEMAS.html
         # https://docs.aws.amazon.com/redshift/latest/dg/r_HAS_TABLE_PRIVILEGE.html
+        schemas_to_exclude = self.configuration.get('exclude_schemas').split(',')
+        full_schemas_to_exclude = '\',\''.join([schema for schema in schemas_to_exclude if '%' not in schema])
+        pattern_schemas_to_exclude = '\' AND table_schema NOT LIKE \''.join([schema for schema in schemas_to_exclude if ' %' in schema])
         query = """
         WITH tables AS (
             SELECT DISTINCT table_name,
                             table_schema,
                             column_name,
                             ordinal_position AS pos
             FROM svv_columns
-            WHERE table_schema NOT IN ('pg_internal','pg_catalog','information_schema')
-            AND table_schema NOT LIKE 'pg_temp_%'
+            WHERE table_schema NOT IN ('{full_schemas_to_exclude}')
+            AND table_schema NOT LIKE '{pattern_schemas_to_exclude}'
         )
         SELECT table_name, table_schema, column_name
         FROM tables
@@ -397,7 +406,10 @@ def _get_tables(self, schema):
                 HAS_TABLE_PRIVILEGE('"' || table_schema || '"."' || table_name || '"', 'SELECT')
             )
         ORDER BY table_name, pos
-        """
+        """.format(
+            full_schemas_to_exclude=full_schemas_to_exclude,
+            pattern_schemas_to_exclude=pattern_schemas_to_exclude
+        )
 
         self._get_definitions(schema, query)
 
@@ -460,6 +472,11 @@ def configuration_schema(cls):
                     "title": "Query Group for Scheduled Queries",
                     "default": "default",
                 },
+                "exclude_schemas": {
+                    "type": "string",
+                    "title": "Exclude schemas from scanning tables",
+                    "default": "pg_internal,pg_catalog,information_schema,pg_temp_%",
+                },
             },
             "order": [
                 "rolename",
@@ -472,6 +489,7 @@ def configuration_schema(cls):
                 "user",
                 "dbname",
                 "sslmode",
+                "exclude_schemas",
                 "adhoc_query_group",
                 "scheduled_query_group",
             ],