Merge pull request #319 from aws-samples/318-row-level-security-enhancement

improve the logic of replacing table name in the original SQL

Author: supinyu

application/nlq/business/datasource/base.py

@@ -1,4 +1,5 @@
 import yaml
+import re
 from abc import ABC, abstractmethod
 
 from nlq.business.login_user import LoginUser
@@ -102,6 +103,18 @@ def replace_table_with_cte(sql, table_config: dict):
             sql_splits.append(sql)
 
         for table_name, sub_query in table_config.items():
+            if '.' in table_name:
+                # 如果表名包含schema name(格式: schema.table), 则将.替换成__
+                schema_name, table_name_alone = table_name.split('.')
+                table_name_replaced = table_name.replace('.', '__')
+                if table_name in sql_splits[1]:
+                    # 替换带schema的表名
+                    sql_splits[1] = re.sub(r'\b{}\b'.format(table_name), table_name_replaced, sql_splits[1])
+                elif table_name_alone in sql_splits[1]:
+                    # 替换不带schema的表名
+                    sql_splits[1] = re.sub(r'\b{}\b'.format(table_name_alone), table_name_replaced, sql_splits[1])
+
+                table_name = table_name_replaced
             cte_sql += f"/* rls applied */ {table_name} AS {sub_query},\n"
         if origin_sql_has_cte:
             cte_sql = cte_sql[:-1]

application/tests/unit_tests/test_row_level_security.py

@@ -10,6 +10,16 @@ def setUp(self):
         self.two_table_join_sql = '''SELECT c.`name`, o.`product`, o.`quantity`, o.`territory`
 FROM customer c
 JOIN orders o ON c.`id` = o.`customer_id`
+LIMIT 100'''
+
+        self.two_table_join_sql_with_schema = '''SELECT c.`name`, o.`product`, o.`quantity`, o.`territory`
+FROM someschema.customer c
+JOIN someschema.orders o ON c.`id` = o.`customer_id`
+LIMIT 100'''
+
+        self.two_table_join_sql_with_schema_output = '''SELECT c.`name`, o.`product`, o.`quantity`, o.`territory`
+FROM someschema__customer c
+JOIN someschema__orders o ON c.`id` = o.`customer_id`
 LIMIT 100'''
 
         self.expected_rls_enabled_sql = (
@@ -18,6 +28,12 @@ def setUp(self):
             "/* rls applied */ orders AS (SELECT * FROM orders WHERE territory = 'Asia')\n"
             f"{self.two_table_join_sql}")
 
+        self.expected_rls_enabled_sql_with_schema = (
+            "WITH\n"
+            "/* rls applied */ someschema__customer AS (SELECT * FROM someschema.customer WHERE created_by = 'admin'),\n"
+            "/* rls applied */ someschema__orders AS (SELECT * FROM someschema.orders WHERE territory = 'Asia')\n"
+            f"{self.two_table_join_sql_with_schema_output}")
+
         self.base = MySQLDataSource()
 
     def test_row_level_security_control(self):
@@ -34,6 +50,24 @@ def test_row_level_security_control(self):
 
         self.assertEqual(self.expected_rls_enabled_sql, rls_modified_sql)
 
+    def test_row_level_security_control_with_schema(self):
+        test_yaml = '''tables:
+    - table_name: someschema.customer
+      columns:
+        - column_name: created_by
+          column_value: $login_user.username
+    - table_name: someschema.orders
+      columns:
+        - column_name: territory
+          column_value: Asia'''
+        rls_modified_sql = self.base.row_level_security_control(self.two_table_join_sql_with_schema, test_yaml, LoginUser('admin'))
+
+        self.assertEqual(self.expected_rls_enabled_sql_with_schema, rls_modified_sql)
+
+        # 测试不带schema的表名的兼容性
+        rls_modified_sql2 = self.base.row_level_security_control(self.two_table_join_sql, test_yaml, LoginUser('admin'))
+        self.assertEqual(self.expected_rls_enabled_sql_with_schema, rls_modified_sql2)
+
     def test_cte_replace1(self):
         original_sql = """SELECT
     offer_id,