使用goInception作为MySQL查询表权限的解析工具

common/templates/config.html

@@ -56,7 +56,7 @@ <h5 class="control-label text-bold">组：</h5>
                         <h4 style="color: darkgrey; display: inline;"><b>Inception配置</b></h4>&nbsp;&nbsp;&nbsp;
                         <button id='check_goincption' class='btn-sm btn-info'>测试goInception连接</button>&nbsp;
                         <button id='check_incption' class='btn-sm btn-info'>测试Inception连接</button>
-                        <h6 style="color:red">注：默认使用goInception进行MySQL的审核和执行，Inception作为查询表权限校验和脱敏解析工具使用</h6>
+                        <h6 style="color:red">注：默认使用goInception进行MySQL的审核执行，Inception作为脱敏解析工具使用</h6>
                         <hr/>
                         <div class="form-horizontal">
                             <div class="form-group">
@@ -278,9 +278,23 @@ <h5 style="color: darkgrey"><b>SQL上线</b></h5>
                             </div>
                         </div>
                         <h5 style="color: darkgrey"><b>SQL查询</b></h5>
-                        <h6 style="color:red">注：开启Inception检测和脱敏功能必须要配置Inception信息，用于SQL语法解析</h6>
+                        <h6 style="color:red">注：开启脱敏功能必须要配置Inception信息，用于SQL语法解析</h6>
                         <hr/>
                         <div class="form-horizontal">
+                            <div class="form-group">
+                                <label for="data_masking"
+                                       class="col-sm-4 control-label">DATA_MASKING</label>
+                                <div class="col-sm-8">
+                                    <div class="switch switch-small">
+                                        <label>
+                                            <input id="data_masking"
+                                                   key="data_masking"
+                                                   value="{{ config.data_masking }}" type="checkbox">
+                                            是否开启动态脱敏
+                                        </label>
+                                    </div>
+                                </div>
+                            </div>
                             <div class="form-group">
                                 <label for="query_check"
                                        class="col-sm-4 control-label">QUERY_CHECK</label>
@@ -290,7 +304,7 @@ <h6 style="color:red">注：开启Inception检测和脱敏功能必须要配置I
                                             <input id="query_check"
                                                    key="query_check"
                                                    value="{{ config.query_check }}" type="checkbox">
-                                            是否开启Inception检测(表权限和脱敏校验)
+                                            是否开启脱敏校验(无法脱敏的语句会抛错)
                                         </label>
                                     </div>
                                 </div>
@@ -309,20 +323,6 @@ <h6 style="color:red">注：开启Inception检测和脱敏功能必须要配置I
                                     </div>
                                 </div>
                             </div>
-                            <div class="form-group">
-                                <label for="data_masking"
-                                       class="col-sm-4 control-label">DATA_MASKING</label>
-                                <div class="col-sm-8">
-                                    <div class="switch switch-small">
-                                        <label>
-                                            <input id="data_masking"
-                                                   key="data_masking"
-                                                   value="{{ config.data_masking }}" type="checkbox">
-                                            是否开启动态脱敏
-                                        </label>
-                                    </div>
-                                </div>
-                            </div>
                             <div class="form-group">
                                 <label for="inception_remote_backup_port"
                                        class="col-sm-4 control-label">MAX_EXECUTION_TIME</label>

sql/engines/goinception.py

@@ -117,6 +117,20 @@ def query(self, db_name=None, sql='', limit_num=0, close_conn=True):
             self.close()
         return result_set
 
+    def query_print(self, instance, db_name=None, sql=''):
+        """
+        打印语法树。
+        """
+        sql = f"""/*--user={instance.user};--password={instance.password};--host={instance.host};--port={instance.port};--enable-query-print;*/
+                          inception_magic_start;\
+                          use `{db_name}`;
+                          {sql}
+                          inception_magic_commit;"""
+        print_info = self.query(db_name=db_name, sql=sql).to_dict()[1]
+        if print_info.get('errmsg'):
+            raise RuntimeError(print_info.get('errmsg'))
+        return print_info
+
     def get_variables(self, variables=None):
         """获取实例参数"""
         if variables:
@@ -140,7 +154,59 @@ def osc_control(self, **kwargs):
             sql = f"inception {command} osc '{sqlsha1}';"
         return self.query(sql=sql)
 
+    @staticmethod
+    def get_table_ref(query_tree, db_name=None):
+        """
+        * 从goInception解析后的语法树里解析出兼容Inception格式的引用表信息。
+        * 目前的逻辑是在SQL语法树中通过递归查找选中最小的 TableRefs 子树（可能有多个），
+        然后在最小的 TableRefs 子树选中Source节点来获取表引用信息。
+        * 查找最小TableRefs子树的方案竟然是通过逐步查找最大子树（直到找不到）来获得的，
+        具体为什么这样实现，我不记得了，只记得当时是通过猜测goInception的语法树生成规
+        则来写代码，结果猜一次错一次错一次猜一次，最终代码逐渐演变于此。或许直接查找最
+        小子树才是效率较高的算法，但是就这样吧，反正它能运行 :)
+        """
+        table_ref = []
+
+        find_queue = [query_tree]
+        for tree in find_queue:
+            tree = DictTree(tree)
+
+            # nodes = tree.find_max_tree("TableRefs") or tree.find_max_tree("Left", "Right")
+            nodes = tree.find_max_tree("TableRefs")
+            if nodes:
+                # assert isinstance(v, dict) is true
+                find_queue.extend([v for node in nodes for v in node.values() if v])
+            else:
+                snodes = tree.find_max_tree("Source")
+                if snodes:
+                    table_ref.extend([
+                        {
+                            "schema": snode['Source']['Schema']['O'] or db_name,
+                            "name": snode['Source']['Name']['O']
+                        } for snode in snodes
+                    ])
+                # assert: source node must exists if table_refs node exists.
+                # else:
+                #     raise Exception("GoInception Error: not found source node")
+        return table_ref
+
     def close(self):
         if self.conn:
             self.conn.close()
             self.conn = None
+
+
+class DictTree(dict):
+    def find_max_tree(self, *keys):
+        """通过广度优先搜索算法查找满足条件的最大子树(不找叶子节点)"""
+        fit = []
+        find_queue = [self]
+        for tree in find_queue:
+            for k, v in tree.items():
+                if k in keys:
+                    fit.append({k: v})
+                elif isinstance(v, dict):
+                    find_queue.append(v)
+                elif isinstance(v, list):
+                    find_queue.extend([n for n in v if isinstance(n, dict)])
+        return fit

sql/engines/mysql.py

@@ -153,6 +153,12 @@ def query_check(self, db_name=None, sql=''):
         if '*' in sql:
             result['has_star'] = True
             result['msg'] = 'SQL语句中含有 * '
+        # select语句先使用Explain判断语法是否正确
+        if re.match(r"^select", sql, re.I):
+            explain_result = self.query(db_name=db_name, sql=f"explain {sql}")
+            if explain_result.error:
+                result['bad_query'] = True
+                result['msg'] = explain_result.error
         return result
 
     def filter_sql(self, sql='', limit_num=0):

sql/query_privileges.py

@@ -22,7 +22,7 @@
 from common.config import SysConfig
 from common.utils.const import WorkflowDict
 from common.utils.extend_json_encoder import ExtendJSONEncoder
-from sql.engines.inception import InceptionEngine
+from sql.engines.goinception import GoInceptionEngine
 from sql.models import QueryPrivilegesApply, QueryPrivileges, Instance, ResourceGroup
 from sql.notify import notify_for_audit
 from sql.utils.resource_group import user_groups, user_instances
@@ -61,30 +61,31 @@ def query_priv_check(user, instance, db_name, sql_content, limit_num):
     # explain和show create跳过权限校验
     if re.match(r"^explain|^show\s+create", sql_content, re.I):
         return result
-    # 其他尝试使用inception解析
-    try:
-        # 尝试使用Inception校验表权限
-        table_ref = _table_ref(f"{sql_content.rstrip(';')};", instance, db_name)
-        # 循环验证权限，可能存在性能问题，但一次查询涉及的库表数量有限，可忽略
-        for table in table_ref:
-            # 既无库权限也无表权限
-            if not _db_priv(user, instance, table['db']) and not _tb_priv(user, instance, db_name, table['table']):
-                result['status'] = 1
-                result['msg'] = f"你无{db_name}.{table['table']}表的查询权限！请先到查询权限管理进行申请"
-                return result
-        # 获取查询涉及库/表权限的最小limit限制，和前端传参作对比，取最小值
-        # 循环获取，可能存在性能问题，但一次查询涉及的库表数量有限，可忽略
-        for table in table_ref:
-            priv_limit = _priv_limit(user, instance, db_name=table['db'], tb_name=table['table'])
-            limit_num = min(priv_limit, limit_num) if limit_num else priv_limit
-        result['data']['limit_num'] = limit_num
-    except SyntaxError as msg:
-        result['status'] = 1
-        result['msg'] = f"SQL语法错误，{msg}"
-        return result
-    except Exception as msg:
-        # 表权限校验失败再次校验库权限
-        # 先获取查询语句涉及的库
+
+    # 仅MySQL做表权限校验
+    if instance.db_type == 'mysql':
+        try:
+            table_ref = _table_ref(f"{sql_content.rstrip(';')};", instance, db_name)
+            # 循环验证权限，可能存在性能问题，但一次查询涉及的库表数量有限
+            for table in table_ref:
+                # 既无库权限也无表权限则鉴权失败
+                if not _db_priv(user, instance, table['schema']) and \
+                        not _tb_priv(user, instance, db_name, table['name']):
+                    result['status'] = 1
+                    result['msg'] = f"你无{table['schema']}.{table['name']}表的查询权限！请先到查询权限管理进行申请"
+                    return result
+            # 获取查询涉及库/表权限的最小limit限制，和前端传参作对比，取最小值
+            for table in table_ref:
+                priv_limit = _priv_limit(user, instance, db_name=table['schema'], tb_name=table['name'])
+                limit_num = min(priv_limit, limit_num) if limit_num else priv_limit
+            result['data']['limit_num'] = limit_num
+        except Exception as msg:
+            logger.error(traceback.format_exc())
+            result['status'] = 1
+            result['msg'] = f"无法校验查询语句权限，请联系管理员，错误信息：{msg}"
+    # 其他类型实例仅校验库权限
+    else:
+        # 先获取查询语句涉及的库，redis、mssql特殊处理，仅校验当前选择的库
         if instance.db_type in ['redis', 'mssql']:
             dbs = [db_name]
         else:
@@ -105,18 +106,6 @@ def query_priv_check(user, instance, db_name, sql_content, limit_num):
             priv_limit = _priv_limit(user, instance, db_name=db_name)
             limit_num = min(priv_limit, limit_num) if limit_num else priv_limit
         result['data']['limit_num'] = limit_num
-
-        # 实例为mysql的，需要判断query_check状态
-        if instance.db_type == 'mysql':
-            # 开启query_check，则禁止执行
-            if SysConfig().get('query_check'):
-                result['status'] = 1
-                result['msg'] = f"无法校验查询语句权限，请检查语法是否正确或联系管理员，错误信息：{msg}"
-                return result
-            # 关闭query_check，标记权限校验为跳过，可继续执行
-            else:
-                result['data']['priv_check'] = False
-
     return result
 
 
@@ -412,19 +401,9 @@ def _table_ref(sql_content, instance, db_name):
     :param db_name:
     :return:
     """
-    if instance.db_type != 'mysql':
-        raise RuntimeError('Inception Error: 仅支持MySQL实例')
-    inception_engine = InceptionEngine()
-    query_tree = inception_engine.query_print(instance=instance, db_name=db_name, sql=sql_content)
-    table_ref = query_tree.get('table_ref', [])
-    db_list = [table_info['db'] for table_info in table_ref]
-    table_list = [table_info['table'] for table_info in table_ref]
-    # 异常解析的情形
-    if '' in db_list or '*' in table_list:
-        raise RuntimeError('Inception Error: 存在空数据库表信息')
-    if not (db_list or table_list):
-        raise RuntimeError('Inception Error: 未解析到任何库表信息')
-    return table_ref
+    engine = GoInceptionEngine()
+    query_tree = engine.query_print(instance=instance, db_name=db_name, sql=sql_content).get('query_tree')
+    return engine.get_table_ref(json.loads(query_tree), db_name=db_name)
 
 
 def _db_priv(user, instance, db_name):