节点在线、应用在线、配置在线使用令牌查询
大石头 authored at 2021-12-16 19:49:30
7.97 KiB
Stardust
/**
 * 数据库自动埋点(pg / mysql2)。
 *
 * 无新增依赖:模块不存在时静默跳过。
 * Span 名称遵循 db:{system}:{action}:{tables} 约定。
 */

'use strict';

const { getTracer } = require('./registry');

/**
 * 截断文本。
 * @param {*} value 原值
 * @param {Number} limit 最大长度
 * @returns {String} 结果
 */
function truncate(value, limit) {
    const text = value === null || value === undefined ? '' : String(value);
    if (limit > 0 && text.length > limit) return text.substring(0, limit);
    return text;
}

/**
 * 尝试加载可选模块。
 * @param {String} name 模块名
 * @returns {Object|null} 模块
 */
function tryRequire(name) {
    try {
        return require(name);
    } catch (err) {
        return null;
    }
}

/**
 * 识别 SQL 动作。
 * @param {String} sql SQL 文本
 * @returns {String} 动作
 */
function detectAction(sql) {
    const text = String(sql || '').trim().toLowerCase();
    if (!text) return 'execute';
    if (/^(select|show|describe|desc|explain|with)\b/.test(text)) return 'select';
    if (/^insert\b/.test(text)) return 'insert';
    if (/^update\b/.test(text)) return 'update';
    if (/^delete\b/.test(text)) return 'delete';
    if (/^replace\b/.test(text)) return 'replace';
    if (/^merge\b/.test(text)) return 'merge';
    if (/^create\b/.test(text)) return 'create';
    if (/^alter\b/.test(text)) return 'alter';
    if (/^drop\b/.test(text)) return 'drop';
    if (/^truncate\b/.test(text)) return 'truncate';
    if (/^(begin|commit|rollback|start\s+transaction)\b/.test(text)) return 'transaction';

    const match = /^([a-z]+)/.exec(text);
    return match ? match[1] : 'execute';
}

/**
 * 去掉引号、架构前缀与日期分表后缀。
 * @param {String} raw 原始表名
 * @returns {String} 表名
 */
function normalizeTable(raw) {
    let name = String(raw || '')
        .replace(/[`"\[\]]/g, '')
        .trim();
    const dot = name.lastIndexOf('.');
    if (dot >= 0) name = name.substring(dot + 1);
    name = name.replace(/[_-]?\d{6,8}$/, '');
    return name;
}

/**
 * 从 SQL 中提取表名。
 * @param {String} sql SQL 文本
 * @returns {String} 逗号分隔的表名,未识别时为空串
 */
function extractTables(sql) {
    const text = String(sql || '');
    const patterns = [
        /\bfrom\s+([`"\[\]\w.$]+)/gi,
        /\binsert\s+into\s+([`"\[\]\w.$]+)/gi,
        /\bupdate\s+([`"\[\]\w.$]+)/gi,
        /\bjoin\s+([`"\[\]\w.$]+)/gi,
        /\bdrop\s+table\s+(?:if\s+exists\s+)?([`"\[\]\w.$]+)/gi,
        /\btruncate\s+table\s+([`"\[\]\w.$]+)/gi,
    ];

    const tables = [];
    for (const pattern of patterns) {
        pattern.lastIndex = 0;
        let match;
        while ((match = pattern.exec(text)) !== null) {
            const name = normalizeTable(match[1]);
            if (name && tables.indexOf(name) === -1) tables.push(name);
        }
    }
    return tables.join(',');
}

/**
 * 构造数据库 Span 名称。
 * @param {String} system 数据库系统
 * @param {String} sql SQL 文本
 * @returns {String} Span 名称
 */
function buildSpanName(system, sql) {
    const action = detectAction(sql);
    const tables = extractTables(sql) || 'unknown';
    return 'db:' + system + ':' + action + ':' + tables;
}

/**
 * 提取调用参数中的 SQL 文本。
 * @param {*} value 第一个参数
 * @returns {String} SQL 文本
 */
function getSqlText(value) {
    if (typeof value === 'string') return value;
    if (value && typeof value === 'object') {
        if (typeof value.text === 'string') return value.text;
        if (typeof value.sql === 'string') return value.sql;
    }
    return '';
}

/**
 * 包装原型方法,处理回调与 Promise 两种调用方式。
 * @param {Object} proto 原型
 * @param {String} method 方法名
 * @param {String} system 数据库系统
 * @returns {void}
 */
function patchQuery(proto, method, system) {
    if (!proto) return;
    const original = proto[method];
    if (typeof original !== 'function' || original.__stardustWrapped) return;

    const wrapped = function () {
        const self = this;
        const tracer = getTracer();
        if (!tracer || (self && self.__stardustDbCall)) return original.apply(self, arguments);

        const args = Array.prototype.slice.call(arguments);
        const sql = getSqlText(args[0]);
        const span = tracer.newSpan(buildSpanName(system, sql));
        span.setTag(truncate(sql, tracer.maxTagLength));

        let settled = false;
        const finish = function (err) {
            if (settled) return;
            settled = true;
            try {
                if (err) span.setError(err);
            } catch (ignore) {
                // 忽略
            }
            try {
                span.finish();
            } catch (ignore) {
                // 忽略
            }
        };

        let callbackIndex = -1;
        for (let i = 1; i < args.length; i++) {
            if (typeof args[i] === 'function') {
                callbackIndex = i;
                break;
            }
        }

        if (callbackIndex >= 0) {
            const callback = args[callbackIndex];
            args[callbackIndex] = function (err) {
                finish(err);
                return callback.apply(this, arguments);
            };
        }

        let result;
        try {
            if (self) self.__stardustDbCall = true;
            result = original.apply(self, args);
        } catch (err) {
            finish(err);
            throw err;
        } finally {
            if (self) self.__stardustDbCall = false;
        }

        if (callbackIndex < 0 && result && typeof result.then === 'function') {
            return result.then(function (value) {
                finish();
                return value;
            }, function (err) {
                finish(err);
                throw err;
            });
        }
        if (callbackIndex < 0) finish();

        return result;
    };

    wrapped.__stardustWrapped = true;
    proto[method] = wrapped;
}

/**
 * 安装 pg 埋点(Client.prototype.query)。
 * @param {Object} tracer 追踪器
 * @param {Object} [pgModule] 模块,缺省时自动加载
 * @returns {Boolean} 是否安装成功
 */
function instrumentPg(tracer, pgModule) {
    const mod = pgModule || tryRequire('pg');
    if (!mod) return false;

    const Client = mod.Client || (mod.default && mod.default.Client);
    if (!Client || !Client.prototype) return false;

    patchQuery(Client.prototype, 'query', 'postgresql');
    return true;
}

/**
 * 解析 mysql2 Connection 构造器。
 * @param {Object} mod 模块
 * @returns {Function|null} 构造器
 */
function resolveMysqlConnection(mod) {
    if (!mod) return null;
    if (mod.Connection && mod.Connection.prototype) return mod.Connection;
    if (mod.default && mod.default.Connection && mod.default.Connection.prototype) return mod.default.Connection;

    const promiseModule = tryRequire('mysql2/promise');
    if (promiseModule && promiseModule.Connection && promiseModule.Connection.prototype) return promiseModule.Connection;
    return null;
}

/**
 * 安装 mysql2 埋点(Connection.prototype.query/execute)。
 * @param {Object} tracer 追踪器
 * @param {Object} [mysqlModule] 模块,缺省时自动加载
 * @returns {Boolean} 是否安装成功
 */
function instrumentMysql2(tracer, mysqlModule) {
    const mod = mysqlModule || tryRequire('mysql2');
    if (!mod) return false;

    const Connection = resolveMysqlConnection(mod);
    if (!Connection) return false;

    patchQuery(Connection.prototype, 'query', 'mysql');
    patchQuery(Connection.prototype, 'execute', 'mysql');
    return true;
}

/**
 * 自动检测并安装数据库埋点。
 * @param {Object} tracer 追踪器
 * @returns {void}
 */
function install(tracer) {
    try {
        instrumentPg(tracer);
    } catch (err) {
        // 忽略
    }
    try {
        instrumentMysql2(tracer);
    } catch (err) {
        // 忽略
    }
}

module.exports = {
    install,
    instrumentPg,
    instrumentMysql2,
    buildSpanName,
    detectAction,
    extractTables,
};