"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
const EventEmitter = require("events");
const core_1 = require("@opencensus/core");
const shimmer = require("shimmer");
class PGPlugin extends core_1.BasePlugin {
    constructor(moduleName) {
        super(moduleName);
        this.internalFileList = {
            '6 - 7': {
                'client': 'lib/client'
            }
        };
    }
    applyPatch() {
        this.logger.debug('Patched PG');
        if (this.internalFilesExports.client) {
            this.logger.debug('patching pq.client.prototype.query');
            shimmer.wrap(this.internalFilesExports.client.prototype, 'query', this.getPatchCreateQuery());
        }
        return this.moduleExports;
    }
    applyUnpatch() {
        shimmer.unwrap(this.internalFilesExports.client.prototype, 'query');
    }
    getPatchCreateQuery() {
        const plugin = this;
        return (original) => {
            return function (...args) {
                const span = plugin.tracer.startChildSpan('pg-query', core_1.SpanKind.CLIENT);
                if (span === null)
                    return original.apply(this, arguments);
                let pgQuery;
                if (arguments.length >= 1) {
                    const args = Array.prototype.slice.call(arguments, 0);
                    plugin.populateLabelsFromInputs(span, args);
                    const callback = args[args.length - 1];
                    if (typeof callback === 'function') {
                        args[args.length - 1] = plugin.patchCallback(callback, span);
                    }
                    else if (typeof args[0] === 'object') {
                        plugin.patchSubmittable(args[0], span);
                    }
                    pgQuery = original.apply(this, args);
                }
                else {
                    pgQuery = original.apply(this, arguments);
                }
                if (pgQuery) {
                    if (pgQuery instanceof EventEmitter) {
                        plugin.tracer.wrapEmitter(pgQuery);
                    }
                    else if (typeof pgQuery.then === 'function') {
                        plugin.patchPromise(pgQuery, span);
                    }
                }
                return pgQuery;
            };
        };
    }
    patchCallback(callback, span) {
        const plugin = this;
        return plugin.tracer.wrap((err, res) => {
            plugin.populateLabelsFromOutputs(span, err, res);
            span.end();
            callback(err, res);
        });
    }
    patchSubmittable(pgQuery, span) {
        const plugin = this;
        let spanEnded = false;
        if (pgQuery.handleError) {
            shimmer.wrap(pgQuery, 'handleError', (origCallback) => {
                return this.tracer.wrap(function (...args) {
                    if (!spanEnded) {
                        const err = args[0];
                        plugin.populateLabelsFromOutputs(span, err);
                        span.end();
                        spanEnded = true;
                    }
                    if (origCallback) {
                        origCallback.apply(this, args);
                    }
                });
            });
        }
        if (pgQuery.handleReadyForQuery) {
            shimmer.wrap(pgQuery, 'handleReadyForQuery', (origCallback) => {
                return this.tracer.wrap(function (...args) {
                    if (!spanEnded) {
                        plugin.populateLabelsFromOutputs(span, null, this._result);
                        span.end();
                        spanEnded = true;
                    }
                    if (origCallback) {
                        origCallback.apply(this, args);
                    }
                });
            });
        }
        return pgQuery;
    }
    patchPromise(promise, span) {
        return promise = promise.then((res) => {
            plugin.populateLabelsFromOutputs(span, null, res);
            span.end();
            return res;
        }, (err) => {
            plugin.populateLabelsFromOutputs(span, err);
            span.end();
            throw err;
        });
    }
    populateLabelsFromInputs(span, args) {
        const queryObj = args[0];
        if (typeof queryObj === 'object') {
            if (queryObj.text) {
                span.addAttribute('query', queryObj.text);
            }
            if (plugin.options.detailedCommands === true && queryObj.values) {
                span.addAttribute('values', queryObj.values);
            }
        }
        else if (typeof queryObj === 'string') {
            span.addAttribute('query', queryObj);
            if (plugin.options.detailedCommands === true && args.length >= 2 && typeof args[1] !== 'function') {
                span.addAttribute('values', args[1]);
            }
        }
    }
    populateLabelsFromOutputs(span, err, res) {
        if (plugin.options.detailedCommands !== true)
            return;
        if (err) {
            span.addAttribute('error', err.toString());
        }
        if (res) {
            span.addAttribute('row_count', res.rowCount);
            span.addAttribute('oid', res.oid);
            span.addAttribute('rows', res.rows);
            span.addAttribute('fields', res.fields);
        }
    }
}
exports.PGPlugin = PGPlugin;
const plugin = new PGPlugin('pg');
exports.plugin = plugin;
//# sourceMappingURL=data:application/json;base64,[encoded data, 6524 characters]