SELECT * FROM t_demo_b WHERE is_del=0

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"

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

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

import pymysql
class Cursor(pymysql.cursors.Cursor):
    def mogrify(self, query: str, args: Union[None, list, dict, tuple] = None) -> str:
        # 在此可以編寫(xiě)處理還合成的SQL邏輯
        mogrify_sql: str = super().mogrify(query, args)
        # 在此可以編寫(xiě)處理合成后的SQL邏輯
        return mogrify_sql
class DictCursor(pymysql.cursors.DictCursorMixin, Cursor):
    """A cursor which returns results as a dictionary"""
    # 直接修改Cursor類(lèi)的`mogrify`方法并不會(huì)影響到`DictCursor`類(lèi)，所以我們也要?jiǎng)?chuàng)建一個(gè)新的`Cursor`類(lèi)。

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

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

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

from contextvars import ContextVar, Token
from typing import Any, Dict, Optional, Set
context: ContextVar[Dict[str, Any]] = ContextVar("context", default={})
class Context(object):
    """基礎(chǔ)的context調(diào)用，支持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):
    """簡(jiǎn)單的處理reset token邏輯，和context管理，只用在業(yè)務(wù)代碼"""
    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

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, ))

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
        # 在此可以編寫(xiě)處理還合成的SQL邏輯
        mogrify_sql: str = super().mogrify(query, args)
        # 在此可以編寫(xiě)處理合成后的SQL邏輯
        return mogrify_sql

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
        # 簡(jiǎn)單示例，實(shí)際上正則的效率會(huì)更好
        for replace_table in replace_table_set:
            if replace_table in query:
                # 替換表名
                query = query.replace(f" {replace_table} ", f" {replace_table}_{tenant_id} ")
                # 替換查詢(xún)條件中帶有表名的
                query = query.replace(f" {replace_table}.", f" {replace_table}_{tenant_id}.")
        mogrify_sql: str = super().mogrify(query, args)
        # 在此可以編寫(xiě)處理合成后的SQL邏輯
        return mogrify_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

SELECT * FROM `case`

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
        # 生成一個(gè)解析完成的SQL對(duì)象
        sql_parse: sql_metadata.Parser = sql_metadata.Parser(query)

        # 新加的一個(gè)屬性，這里存下需要校驗(yàn)查詢(xún)條件的表名
        check_flag = False 
        where_table_set: Set[str] = context.where_table_set
        # 該方法會(huì)獲取到SQL對(duì)應(yīng)的table，返回的是一個(gè)table的數(shù)組
        for table_name in sql_parse.tables:
            if table_name in where_table_set:
                if sql_parse.columns_dict:
                    # 該方法會(huì)返回SQL對(duì)應(yīng)的字段，其中分為select, join, where等，這里只用到了where
                    for where_column in sql_parse.columns_dict.get("where", []):
                        # 如果連表，里面存的是類(lèi)似于t_demo.tenant_id，所以要兼容這一個(gè)情況
                        if "tenant_id" in where_column.lower().split("."):
                            check_flag = True
                            break
        if not check_flag:
            # 檢查不通過(guò)就拋錯(cuò)
            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存放著解析完的數(shù)據(jù)，比如SELECT * FROM t_demo解析后是
                # [SELECT, *, FROM, t_demo]四個(gè)token
                for token in sql_parse.tokens:
                    # 判斷token是否是表名  
                    if token.is_potential_table_name:
                        # 提取規(guī)范的表名
                        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的下一個(gè)字段
                            if token.next_token.normalized != "AS":
                                # 如果當(dāng)前表沒(méi)有設(shè)置別名
                                # 通過(guò)AS把替換前的表名設(shè)置為新表名的別名，這樣一來(lái)后面的表名即使沒(méi)進(jìn)行更改，也是能讀到對(duì)應(yīng)商戶ID的表
                                new_table_name += f" AS {parse_table_name}"
                            query += new_table_name
                            continue
                    # 通過(guò)stringified_token獲取的數(shù)據(jù)會(huì)自動(dòng)帶空格，比如`FROM`得到的會(huì)是` FROM`，這樣拼接的時(shí)候就不用考慮是否加空格了
                    new_query += token.stringified_token
        mogrify_sql: str = super().mogrify(new_query, args)
        # 在此可以編寫(xiě)處理合成后的SQL邏輯
        return mogrify_sql

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 {}