Python探針怎麼完成呼叫函式庫的資料擷取

1.簡單粗暴的方法--對mysql函式庫進行封裝

要統計一個執行過程， 就需要知道這個執行過程的開始位置和結束位置， 所以最簡單粗暴的方法就是基於要呼叫的方法進行封裝，在框架呼叫MySQL函式庫和MySQL函式庫中間實作一個中間層, 在中間層完成耗時統計,如:

# 伪代码
def my_execute(conn, sql, param):
 # 针对MySql库的统计封装组件
 with MyTracer(conn, sql, param):
     # 以下为正常使用MySql库的代码
with conn.cursor as cursor:
 cursor.execute(sql, param)
...

看樣子實作起來非常好,而且更改非常方便, 但由於是在最頂層的API上進行修改, 其實是非常不靈活的, 同時在cursor.execute裡會進行一些預操作, 如把sql和param進行拼接, 調用nextset清除當前遊標的數據等等。我們最後拿到的數據如時間耗時也是不準確的, 同時也沒辦法得到一些詳細的元數據， 如錯誤碼等等.

如果要拿到最直接有用的數據，就只能去改源代碼, 然後再調用源代碼了, 但是如果每個庫都需要改源代碼才能統計, 那也太麻煩了, 好在Python也提供了一些類似探針的接口, 可以通過探針把庫的源碼進行替換完成我們的程式碼.

2.Python的探針

在Python中可以透過sys.meta_path來實現import hook的功能, 當執行import 相關操作時,會根據sys.meta_path定義的物件對import相關函式庫進行更改.sys.meta_path中的物件需要實作一個find_module方法， 這個find_module方法傳回None或一個實作了load_module方法的物件, 我們可以透過這個物件, 針對一些函式庫在import時, 把相關的方法進行替換, 簡單用法如下,通過hooktime.sleep讓他在sleep的時候能打印消耗的時間.

import importlib
import sys
from functools import wraps
def func_wrapper(func):
    """这里通过一个装饰器来达到狸猫换太子和获取数据的效果"""
    @wraps(func)
    def wrapper(*args, **kwargs):
        # 记录开始时间
        start = time.time()
        result = func(*args, **kwargs)
        # 统计消耗时间
        end = time.time()
        print(f"speed time:{end - start}")
        return result
    return wrapper
class MetaPathFinder:
    def find_module(self, fullname, path=None):
        # 执行时可以看出来在import哪些模块
        print(f'find module:{path}:{fullname}')
        return MetaPathLoader()
class MetaPathLoader:
    def load_module(self, fullname):
        # import的模块都会存放在sys.modules里面， 通过判断可以减少重复import
        if fullname in sys.modules:
            return sys.modules[fullname]
        # 防止递归调用
        finder = sys.meta_path.pop(0)
        # 导入 module
        module = importlib.import_module(fullname)
        if fullname == 'time':
            # 替换函数
            module.sleep = func_wrapper(module.sleep)
        sys.meta_path.insert(0, finder)
        return module
sys.meta_path.insert(0, MetaPathFinder())
if __name__ == '__main__':
    import time
    time.sleep(1)
# 输出示例:
# find module:datetime
# find module:time
# load module:time
# find module:math
# find module:_datetime
# speed time:1.00073385238647468

3.製作探針模組

#了解完了主要流程後， 可以開始製作自己的探針模組了, 由於示例只涉及到aiomysql模組, 那麼在MetaPathFinder.find_module中需要只對aiomysql模組進行處理, 其他的先忽略. 然後我們需要確定我們要把aiomysql的哪個功能給替換, 從業務上來說, 一般情況下我們只要cursor.execute, cursor.fetchone, cursor.fetchall, cursor.executemany這幾個主要的操作,所以需要深入cursor看看如何去更改代碼,後者重載哪個函數.

先cursor.execute的源碼(cursor.executemanay也類似), 發現會先調用self.nextset的方法, 把上個請求的資料先拿完, 再合併sql語句, 最後透過self._query進行查詢:

async def execute(self, query, args=None):
    """Executes the given operation
    Executes the given operation substituting any markers with
    the given parameters.
    For example, getting all rows where id is 5:
        cursor.execute("SELECT * FROM t1 WHERE id = %s", (5,))
    :param query: ``str`` sql statement
    :param args: ``tuple`` or ``list`` of arguments for sql query
    :returns: ``int``, number of rows that has been produced of affected
    """
    conn = self._get_db()

    while (await self.nextset()):
        pass

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

    await self._query(query)
    self._executed = query
    if self._echo:
        logger.info(query)
        logger.info("%r", args)
    return self._rowcount

再看cursor.fetchone的原始碼(cursor.fetchall也類似), 發現其實是從快取中取得資料,

#這些資料在執行cursor.execute中就已經取得了:

def fetchone(self):
    """Fetch the next row """
    self._check_executed()
    fut = self._loop.create_future()
    if self._rows is None or self._rownumber >= len(self._rows):
        fut.set_result(None)
        return fut
    result = self._rows[self._rownumber]
    self._rownumber += 1
    fut = self._loop.create_future()
    fut.set_result(result)
    return fut

綜合上面的分析, 我們只要對核心的方法self._query進行重載即可拿到我們要的資料, 從源碼中我們可以知道, 我們能獲取到傳入self._query的self和sql參數, 根據self又能獲取到查詢的結果, 同時我們透過裝飾器能獲取到運行的時間, 要的數據基本上都到齊了,

依照思路修改後的程式碼如下：

import importlib
import time
import sys
from functools import wraps
from typing import cast, Any, Callable, Optional, Tuple, TYPE_CHECKING
from types import ModuleType
if TYPE_CHECKING:
    import aiomysql
def func_wrapper(func: Callable):
    @wraps(func)
    async def wrapper(*args, **kwargs) -> Any:
        start: float = time.time()
        func_result: Any = await func(*args, **kwargs)
        end: float = time.time()
        # 根据_query可以知道, 第一格参数是self, 第二个参数是sql
        self: aiomysql.Cursor = args[0]
        sql: str = args[1]
        # 通过self,我们可以拿到其他的数据
        db: str = self._connection.db
        user: str = self._connection.user
        host: str = self._connection.host
        port: str = self._connection.port
        execute_result: Tuple[Tuple] = self._rows
        # 可以根据自己定义的agent把数据发送到指定的平台, 然后我们就可以在平台上看到对应的数据或进行监控了, 
        # 这里只是打印一部分数据出来
        print({
            "sql": sql,
            "db": db,
            "user": user,
            "host": host,
            "port": port,
            "result": execute_result,
            "speed time": end - start
        })
        return func_result
    return cast(Callable, wrapper)
class MetaPathFinder:

    @staticmethod
    def find_module(fullname: str, path: Optional[str] = None) -> Optional["MetaPathLoader"]:
        if fullname == 'aiomysql':
            # 只有aiomysql才进行hook
            return MetaPathLoader()
        else:
            return None
class MetaPathLoader:
    @staticmethod
    def load_module(fullname: str):
        if fullname in sys.modules:
            return sys.modules[fullname]
        # 防止递归调用
        finder: "MetaPathFinder" = sys.meta_path.pop(0)
        # 导入 module
        module: ModuleType = importlib.import_module(fullname)
        # 针对_query进行hook
        module.Cursor._query = func_wrapper(module.Cursor._query)
        sys.meta_path.insert(0, finder)
        return module
async def test_mysql() -> None:
    import aiomysql
    pool: aiomysql.Pool = await aiomysql.create_pool(
        host='127.0.0.1', port=3306, user='root', password='123123', db='mysql'
    )
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute("SELECT 42;")
            (r,) = await cur.fetchone()
            assert r == 42
    pool.close()
    await pool.wait_closed()

if __name__ == '__main__':
    sys.meta_path.insert(0, MetaPathFinder())
    import asyncio
    asyncio.run(test_mysql())
# 输出示例:
# 可以看出sql语句与我们输入的一样, db, user, host, port等参数也是, 还能知道执行的结果和运行时间
# {'sql': 'SELECT 42;', 'db': 'mysql', 'user': 'root', 'host': '127.0.0.1', 'port': 3306, 'result': ((42,),), 'speed time': 0.00045609474182128906}

這個範例看來很不錯, 但是需要在呼叫的入口處顯式呼叫該邏輯, 通常一個專案可能有幾個入口, 每個入口都顯示調用該邏輯會非常麻煩， 而且必須先調用我們的hook邏輯後才能import， 這樣就得訂好引入規範， 不然就可能出現部分地方hook不成功,如果能把引入hook這個邏輯安排在解析器啟動後馬上執行， 就可以完美地解決這個問題了. 查閱了一翻資料後發現,python解釋器初始化的時候會自動import PYTHONPATH下存在的sitecustomize和usercustomize模組, 我們只要創建該模組, 並在模組裡面寫入我們的替換函數即可。

.
├── __init__.py
├── hook_aiomysql.py
├── sitecustomize.py
└── test_auto_hook.py

hook_aiomysql.py是我們製作探針的程式碼為例子， 而sitecustomize.py存放的程式碼如下, 非常簡單， 就是引入我們的探針程式碼, 並插入到sys.meta_path：

import sys
from hook_aiomysql import MetaPathFinder
sys.meta_path.insert(0, MetaPathFinder())

test_auto_hook.py則是測試程式碼：

import asyncio
from hook_aiomysql import test_mysql
asyncio.run(test_mysql())

接下來只要設定PYTHONPATH並執行我們的程式碼即可（如果是專案的話一般交由superisor啟動，則可以在設定檔中設定好PYTHONPATH）：

(.venv) ➜  python_hook git:(master) ✗ export PYTHONPATH=.      
(.venv) ➜  python_hook git:(master) ✗ python test_auto_hook.py 
{'sql': 'SELECT 42;', 'db': 'mysql', 'user': 'root', 'host': '127.0.0.1', 'port': 3306, 'result': ((42,),), 'speed time': 0.000213623046875}

4.直接取代方法

#可以看到上面的方法很好的運作了， 而且可以很方便的嵌入到我們的專案中， 但是依賴與sitecustomize.py檔案很難讓他抽離成一個第三方的函式庫， 如果要抽離成第三方的函式庫就得考慮看看有沒有其他的方法。在上面介紹MetaPathLoader時說到了sys.module, 在裡面透過sys.modules來減少重複引入：

class MetaPathLoader:
    def load_module(self, fullname):
        # import的模块都会存放在sys.modules里面， 通过判断可以减少重复import
        if fullname in sys.modules:
            return sys.modules[fullname]
        # 防止递归调用
        finder = sys.meta_path.pop(0)
        # 导入 module
        module = importlib.import_module(fullname)
        if fullname == 'time':
            # 替换函数
            module.sleep = func_wrapper(module.sleep)
        sys.meta_path.insert(0, finder)
        return module

這個減少重複引入的原理是， 每次引入一個模組後， 他就會存放在sys .modules， 如果是重複引入， 就會直接刷新成最新引入的模組。上面之所以會考慮到減少重複import是因為我們不會在程式運行時升級第三方函式庫的依賴。利用到我們可以不考慮重複引入同名不同實現的模組， 以及sys.modules會緩存引入模組的特點， 我們可以把上面的邏輯簡化成引入模組->替換當前模組方法為我們修改的hook方法。

import time
from functools import wraps
from typing import Any, Callable, Tuple, cast
import aiomysql
def func_wrapper(func: Callable):
    """和上面一样的封装函数, 这里简单略过"""
# 判断是否hook过
_IS_HOOK: bool = False
# 存放原来的_query
_query: Callable = aiomysql.Cursor._query
# hook函数
def install_hook() -> None:
    _IS_HOOK = False
    if _IS_HOOK:
        return
    aiomysql.Cursor._query = func_wrapper(aiomysql.Cursor._query)
    _IS_HOOK = True
# 还原到原来的函数方法
def reset_hook() -> None:
    aiomysql.Cursor._query = _query
    _IS_HOOK = False

代码简单明了，接下来跑一跑刚才的测试：

import asyncio
import aiomysql
from demo import install_hook, reset_hook
async def test_mysql() -> None:
    pool: aiomysql.Pool = await aiomysql.create_pool(
        host='127.0.0.1', port=3306, user='root', password='', db='mysql'
    )
    async with pool.acquire() as conn:
        async with conn.cursor() as cur:
            await cur.execute("SELECT 42;")
            (r,) = await cur.fetchone()
            assert r == 42
    pool.close()
    await pool.wait_closed()

print("install hook")
install_hook()
asyncio.run(test_mysql())
print("reset hook")
reset_hook()
asyncio.run(test_mysql())
print("end")

通过测试输出可以发现我们的逻辑的正确的, install hook后能出现我们提取的元信息， 而reset后则不会打印原信息

install hook
{'sql': 'SELECT 42;', 'db': 'mysql', 'user': 'root', 'host': '127.0.0.1', 'port': 3306, 'result': ((42,),), 'speed time': 0.000347137451171875}
reset hook
end

以上是Python探針怎麼完成呼叫函式庫的資料擷取的詳細內容。更多資訊請關注PHP中文網其他相關文章！