How to modify business SQL code based on runtime in Python?

1. Origin

Recently, the project is preparing to implement SASS. One of the characteristics of SASS is multi-tenancy, and the data between each tenant must be isolated. Common database isolation solutions include databases. Isolation, table isolation, field isolation, currently I only use table isolation and field isolation (the principles of database isolation are similar). Field isolation is relatively simple, that is, the query conditions are different. For example, the following SQL query:

SELECT * FROM t_demo WHERE tenant_id='xxx' AND is_del=0

But for the sake of rigor, it is necessary to check whether the corresponding table is included before executing the SQL The query field of tenant_id.

For table isolation, it is a bit more troublesome. It needs to process a certain data table according to the corresponding tenant ID during operation. For example, if there is a SQL query like the following:

SELECT * FROM t_demo WHERE is_del=0

When encountering tenant A, the SQL query will become:

SELECT * FROM t_demo_a WHERE is_del=0

When encountering tenant B, the SQL query will become:

SELECT * FROM t_demo_b WHERE is_del=0

If the number of merchants is fixed, it is usually enough to write if-else in the code to judge. However, merchants for common SASS applications will be added all the time, so for this The SQL logic will become like this:

def sql_handle(tenant_id: str):
    table_name: str = f"t_demo_{tenant_id}"
    sql: str = f"SELECT * FROM {table_name} WHERE is_del=0"

But there are several problems with this. For ORM, it is enough to only create a table object corresponding to t_demo at the beginning, but now it has to It is unrealistic to create multiple table objects based on multiple merchants. Secondly, if you write naked SQL, you will generally use IDE to check it. For such SQL:

sql: str = f"SELECT * FROM {table_name} WHERE is_del=0"

IDE cannot check it. , of course, there is another most serious problem, that is, the current project is already very large. If each related table call is adapted and changed, the amount of work will be very huge, so the best solution is to use the engine library After getting the SQL statement passed by the user and before sending it to the MySQL server, the SQL is automatically changed according to the merchant ID. To achieve this effect, we must invade the MySQL## we use. # engine library, modify the methods inside to be compatible with our needs.

Whether you use

dbutils or sqlalchemy, you can specify an engine library. The currently commonly used engine library is pymysql, so The following will be explained using pymysql as an example.

2. Invade the library

Since we must invade the engine library we use, we should first determine which method of the engine library we need to modify. After reading the source code, I It is determined that only the

mogrify method of pymysql.cursors.Cursor needs to be changed:

def mogrify(self, query, args=None):
    """
    Returns the exact string that is sent to the database by calling the
    execute() method.

    This method follows the extension to the DB API 2.0 followed by Psycopg.
    """
    conn = self._get_db()

    if args is not None:
        query = query % self._escape_args(args, conn)
    return query

The function of this method is to integrate the SQL and parameters passed by the user to generate a final The SQL just meets our needs, so we can create a new

Cursor class of our own through inheritance:

import pymysql
class Cursor(pymysql.cursors.Cursor):
    def mogrify(self, query: str, args: Union[None, list, dict, tuple] = None) -> str:
        # 在此可以编写处理还合成的SQL逻辑
        mogrify_sql: str = super().mogrify(query, args)
        # 在此可以编写处理合成后的SQL逻辑
        return mogrify_sql
class DictCursor(pymysql.cursors.DictCursorMixin, Cursor):
    """A cursor which returns results as a dictionary"""
    # 直接修改Cursor类的`mogrify`方法并不会影响到`DictCursor`类，所以我们也要创建一个新的`Cursor`类。

Created

Cursor After the class, we need to consider how to apply our custom Cursor class in pymysql. The general Mysql connection library supports us passing in custom Cursor class, such as pymysql:

import pymysql.cursors
# Connect to the database
connection = pymysql.connect(
    host='localhost',
    user='user',
    password='passwd',
    database='db',
    charset='utf8mb4',
    cursorclass=pymysql.cursors.DictCursor
)

We can specify our

Cursor class through cursorclass, if using If the library does not support it or for other reasons, you need to use the monkey patch method. For specific usage methods, see Python probe to complete the data extraction of the calling library.

3. Get the merchant ID

Now that we have figured out where to modify the SQL, we need to think about how to get the merchant ID in the

mogrify method And those tables need to be replaced. Generally, when we call a piece of code, there are two ways to pass parameters. One is to pass array type parameters:

with conn.cursor() as cursor:
    cursor.execute("SELECT * FROM t_demo WHERE is_del=%s", (0, ))

The other is to pass dictionary type parameters:

with conn.cursor() as cursor:
    cursor.execute("SELECT * FROM t_demo WHERE is_del=%(is_del)s", {"is_del": 0})

Most projects currently have these two types of writing habits, and the engine library will process the parameters

sql and # when executing execute ##args is passed to mogrify. If we use dictionary type parameters, we can embed the parameters we need in it and extract them in mogrify, but If you use array type parameters or an ORM library, it will be more difficult to pass parameters to the mogrify method. In this case, you can pass the parameters to mogrify# implicitly through context. ##Method, specific analysis and principles can be found in: How to use contextvars module source code analysis in python. The usage of context

is very simple. First, create a

context encapsulated class:

from contextvars import ContextVar, Token
from typing import Any, Dict, Optional, Set
context: ContextVar[Dict[str, Any]] = ContextVar("context", default={})
class Context(object):
    """基础的context调用，支持Type Hints检查"""
    tenant_id: str
    replace_table_set: Set[str]
    def __getattr__(self, key: str) -> Any:
        value: Any = context.get().get(key)
        return value
    def __setattr__(self, key: str, value: Any) -> None:
        context.get()[key] = value
class WithContext(Context):
    """简单的处理reset token逻辑，和context管理，只用在业务代码"""
    def __init__(self) -> None:
        self._token: Optional[Token] = None

    def __enter__(self) -> "WithContext":
        self._token = context.set({})
        return self
    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
        if self._token:
            context.reset(self._token)
            self._token = None

Next, in the business code, pass context Pass in the parameters corresponding to the current business:

with WithContext as context:
    context.tenant_id = "xxx"
    context.replace_table_set = {"t_demo"}
    with conn.cursor() as cursor:
        cursor.execute("SELECT * FROM t_demo WHERE is_del=%s", (0, ))

and then call

context

in

mogrify to get the corresponding parameters:

import pymysql
class Cursor(pymysql.cursors.Cursor):
    def mogrify(self, query: str, args: Union[None, list, dict, tuple] = None) -> str:
        tenant_id: str = context.tenant_id
        replace_table_set: Set[str] = context.replace_table_set
        # 在此可以编写处理还合成的SQL逻辑
        mogrify_sql: str = super().mogrify(query, args)
        # 在此可以编写处理合成后的SQL逻辑
        return mogrify_sql

4 .Modify SQL

Now, everything is ready, only the logic of modifying SQL is left. When working on other projects before, the tables built were very standardized, and they started with

t_xxx## Name the table in the format of #, which makes it very convenient to replace the table name. It only needs to be replaced twice to be compatible with most situations. The code is as follows:

import pymysql
class Cursor(pymysql.cursors.Cursor):
    def mogrify(self, query: str, args: Union[None, list, dict, tuple] = None) -> str:
        tenant_id: str = context.tenant_id
        replace_table_set: Set[str] = context.replace_table_set
        # 简单示例，实际上正则的效率会更好
        for replace_table in replace_table_set:
            if replace_table in query:
                # 替换表名
                query = query.replace(f" {replace_table} ", f" {replace_table}_{tenant_id} ")
                # 替换查询条件中带有表名的
                query = query.replace(f" {replace_table}.", f" {replace_table}_{tenant_id}.")
        mogrify_sql: str = super().mogrify(query, args)
        # 在此可以编写处理合成后的SQL逻辑
        return mogrify_sql

但是现在项目的SQL规范并不是很好，有些表名还是MySQL的关键字，所以靠简单的替换是行不通的，同时这个需求中，一些表只需要字段隔离，需要确保有带上对应的字段查询，这就意味着必须有一个库可以来解析SQL，并返回一些数据使我们可以比较方便的知道SQL中哪些是表名，哪些是查询字段了。

目前在Python中有一个比较知名的SQL解析库--sqlparse，它可以通过解析引擎把SQL解析成一个Python对象,之后我们就可以通过一些语法来判断哪些是SQL关键字， 哪些是表名，哪些是查询条件等等。但是这个库只实现一些底层的API，我们需要对他和SQL比较了解之后才能实现一些比较完备的功能，比如下面3种常见的SQL:

SELECT * FROM t_demo
SELECT * FROM t_demo as demo
SELECT * FROM t_other as other LEFT JOIN t_demo demo on demo.xxx==other.xxx

如果我们要通过sqlparse来提取表名的话就需要处理这3种情况，而我们如果要每一个情况都编写出来的话，那将会非常费心费力，同时也可能存在遗漏的情况，这时就需要用到另外一个库--sql_metadata，这个库是基于sqlparse和正则的解析库，同时提供了大量的常见使用方法的封装，我们通过直接调用对应的函数就能知道SQL中有哪些表名，查询字段是什么了。

目前已知这个库有一个缺陷，就是会自动去掉字段的符号， 比如表名为关键字时，我们需要使用`符号把它包起来：

SELECT * FROM `case`

但在经过sql_metadata解析后得到的表名是case而不是`case`，需要人为的处理，但是我并不觉得这是一个BUG，自己不按规范创建表，能怪谁呢。

接下来就可以通过sql_metadata的方法来实现我需要的功能了，在根据需求修改后，代码长这样(说明见注释)：

from typing import Dict, Set, Tuple, Union
import pymysql
import sql_metadata
class Cursor(pymysql.cursors.Cursor):
    def mogrify(self, query: str, args: Union[None, list, dict, tuple] = None) -> str:
        tenant_id: str = context.tenant_id
        # 生成一个解析完成的SQL对象
        sql_parse: sql_metadata.Parser = sql_metadata.Parser(query)

        # 新加的一个属性，这里存下需要校验查询条件的表名
        check_flag = False 
        where_table_set: Set[str] = context.where_table_set
        # 该方法会获取到SQL对应的table，返回的是一个table的数组
        for table_name in sql_parse.tables:
            if table_name in where_table_set:
                if sql_parse.columns_dict:
                    # 该方法会返回SQL对应的字段，其中分为select, join, where等，这里只用到了where
                    for where_column in sql_parse.columns_dict.get("where", []):
                        # 如果连表，里面存的是类似于t_demo.tenant_id，所以要兼容这一个情况
                        if "tenant_id" in where_column.lower().split("."):
                            check_flag = True
                            break
        if not check_flag:
            # 检查不通过就抛错
            raise RuntimeError()
        # 更换表名的逻辑
        replace_table_set: Set[str] = context.replace_table_set
        new_query: str = query
        for table_name in sql_parse.tables:
            if table_name in replace_table_set:
                new_query = ""
                # tokens存放着解析完的数据，比如SELECT * FROM t_demo解析后是
                # [SELECT, *, FROM, t_demo]四个token
                for token in sql_parse.tokens:
                    # 判断token是否是表名  
                    if token.is_potential_table_name:
                        # 提取规范的表名
                        parse_table_name: str = token.stringified_token.strip()
                        if parse_table_name in replace_table_set:
                            new_table_name: str = f" {parse_table_name}_{tenant_id}"
                            # next_token代表SQL的下一个字段
                            if token.next_token.normalized != "AS":
                                # 如果当前表没有设置别名
                                # 通过AS把替换前的表名设置为新表名的别名，这样一来后面的表名即使没进行更改，也是能读到对应商户ID的表
                                new_table_name += f" AS {parse_table_name}"
                            query += new_table_name
                            continue
                    # 通过stringified_token获取的数据会自动带空格，比如`FROM`得到的会是` FROM`，这样拼接的时候就不用考虑是否加空格了
                    new_query += token.stringified_token
        mogrify_sql: str = super().mogrify(new_query, args)
        # 在此可以编写处理合成后的SQL逻辑
        return mogrify_sql

这份代码十分简单，它只做简单介绍，事实上这段逻辑会应用到所有的SQL查询中，我们应该要保证这段代码是没问题的，同时不要有太多的性能浪费，所以在使用的时候要考虑到代码拆分和优化。 比如在使用的过程中可以发现，我们的SQL转换和检查都是在父类的Cursor.mogrify之前进行的，这就意味着不管我们代码逻辑里cursor.execute传的参数是什么，对于同一个代码逻辑来说，传过来的query值是保持不变的，比如下面的代码:

def get_user_info(uid: str) -> Dict[str, Any]:
    with conn.cursor() as cursor:
        cursor.execute("SELECT * FROM t_user WHERE uid=%(uid)s", {"uid": uid})
        return cursor.fetchone() or {}

这段代码中传到Cursor.mogrify的query永远为SELECT * FROM t_user WHERE uid=%(uid)s，有变化的只是args中uid的不同。 有了这样的一个前提条件，那么我们就可以把query的校验结果和转换结果缓存下来，减少每次都需要解析SQL再校验造成的性能浪费。至于如何实现缓存则需要根据自己的项目来决定，比如项目中只有几百个SQL执行，那么直接用Python的dict来存放就可以了，如果项目中执行的SQL很多，同时有些执行的频率非常的高，有些执行的频率非常的低，那么可以考虑使用LRU来缓存。

The above is the detailed content of How to modify business SQL code based on runtime in Python?. For more information, please follow other related articles on the PHP Chinese website!