chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,114 @@
|
||||
const mssql = require("mssql");
|
||||
const { ConnectionStringParser } = require("./utils");
|
||||
|
||||
class MSSQLConnector {
|
||||
#connected = false;
|
||||
database_id = "";
|
||||
connectionConfig = {
|
||||
user: null,
|
||||
password: null,
|
||||
database: null,
|
||||
server: null,
|
||||
port: null,
|
||||
pool: {
|
||||
max: 10,
|
||||
min: 0,
|
||||
idleTimeoutMillis: 30000,
|
||||
},
|
||||
options: {
|
||||
encrypt: false,
|
||||
trustServerCertificate: true,
|
||||
},
|
||||
};
|
||||
|
||||
constructor(
|
||||
config = {
|
||||
// we will force into RFC-3986 from DB
|
||||
// eg: mssql://user:password@server:port/database?{...opts}
|
||||
connectionString: null, // we will force into RFC-3986
|
||||
}
|
||||
) {
|
||||
this.className = "MSSQLConnector";
|
||||
this.connectionString = config.connectionString;
|
||||
this._client = null;
|
||||
this.#parseDatabase();
|
||||
}
|
||||
|
||||
#parseDatabase() {
|
||||
const connectionParser = new ConnectionStringParser({ scheme: "mssql" });
|
||||
const parsed = connectionParser.parse(this.connectionString);
|
||||
|
||||
this.database_id = parsed?.endpoint;
|
||||
this.connectionConfig = {
|
||||
...this.connectionConfig,
|
||||
user: parsed?.username,
|
||||
password: parsed?.password,
|
||||
database: parsed?.endpoint,
|
||||
server: parsed?.hosts?.[0]?.host,
|
||||
port: parsed?.hosts?.[0]?.port,
|
||||
options: {
|
||||
...this.connectionConfig.options,
|
||||
encrypt: parsed?.options?.encrypt === "true",
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
async connect() {
|
||||
this._client = await mssql.connect(this.connectionConfig);
|
||||
this.#connected = true;
|
||||
return this._client;
|
||||
}
|
||||
|
||||
/**
|
||||
*
|
||||
* @param {string} queryString the SQL query to be run
|
||||
* @param {Array} params optional parameters for prepared statement
|
||||
* @returns {Promise<import(".").QueryResult>}
|
||||
*/
|
||||
async runQuery(queryString = "", params = []) {
|
||||
const result = { rows: [], count: 0, error: null };
|
||||
try {
|
||||
if (!this.#connected) await this.connect();
|
||||
|
||||
const request = this._client.request();
|
||||
params.forEach((value, index) => {
|
||||
request.input(`p${index}`, value);
|
||||
});
|
||||
const query = await request.query(queryString);
|
||||
result.rows = query.recordset;
|
||||
result.count = query.rowsAffected.reduce((sum, a) => sum + a, 0);
|
||||
} catch (err) {
|
||||
console.log(this.className, err);
|
||||
result.error = err.message;
|
||||
} finally {
|
||||
// Check client is connected before closing since we use this for validation
|
||||
if (this._client) {
|
||||
await this._client.close();
|
||||
this.#connected = false;
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
async validateConnection() {
|
||||
try {
|
||||
const result = await this.runQuery("SELECT 1");
|
||||
return { success: !result.error, error: result.error };
|
||||
} catch (error) {
|
||||
return { success: false, error: error.message };
|
||||
}
|
||||
}
|
||||
|
||||
getTablesSql() {
|
||||
return `SELECT name FROM sysobjects WHERE xtype='U';`;
|
||||
}
|
||||
|
||||
getTableSchemaSql(table_name) {
|
||||
return {
|
||||
query: `SELECT COLUMN_NAME, COLUMN_DEFAULT, IS_NULLABLE, DATA_TYPE FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = @p0`,
|
||||
params: [table_name],
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
module.exports.MSSQLConnector = MSSQLConnector;
|
||||
@@ -0,0 +1,83 @@
|
||||
const mysql = require("mysql2/promise");
|
||||
const { ConnectionStringParser } = require("./utils");
|
||||
|
||||
class MySQLConnector {
|
||||
#connected = false;
|
||||
database_id = "";
|
||||
constructor(
|
||||
config = {
|
||||
connectionString: null,
|
||||
}
|
||||
) {
|
||||
this.className = "MySQLConnector";
|
||||
this.connectionString = config.connectionString;
|
||||
this._client = null;
|
||||
this.database_id = this.#parseDatabase();
|
||||
}
|
||||
|
||||
#parseDatabase() {
|
||||
const connectionParser = new ConnectionStringParser({ scheme: "mysql" });
|
||||
const parsed = connectionParser.parse(this.connectionString);
|
||||
return parsed?.endpoint;
|
||||
}
|
||||
|
||||
async connect() {
|
||||
this._client = await mysql.createConnection({ uri: this.connectionString });
|
||||
this.#connected = true;
|
||||
return this._client;
|
||||
}
|
||||
|
||||
/**
|
||||
*
|
||||
* @param {string} queryString the SQL query to be run
|
||||
* @param {Array} params optional parameters for prepared statement
|
||||
* @returns {Promise<import(".").QueryResult>}
|
||||
*/
|
||||
async runQuery(queryString = "", params = []) {
|
||||
const result = { rows: [], count: 0, error: null };
|
||||
try {
|
||||
if (!this.#connected) await this.connect();
|
||||
const [query] =
|
||||
params.length > 0
|
||||
? await this._client.execute(queryString, params)
|
||||
: await this._client.query(queryString);
|
||||
result.rows = query;
|
||||
result.count = query?.length;
|
||||
} catch (err) {
|
||||
console.log(this.className, err);
|
||||
result.error = err.message;
|
||||
} finally {
|
||||
// Check client is connected before closing since we use this for validation
|
||||
if (this._client) {
|
||||
await this._client.end();
|
||||
this.#connected = false;
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
async validateConnection() {
|
||||
try {
|
||||
const result = await this.runQuery("SELECT 1");
|
||||
return { success: !result.error, error: result.error };
|
||||
} catch (error) {
|
||||
return { success: false, error: error.message };
|
||||
}
|
||||
}
|
||||
|
||||
getTablesSql() {
|
||||
return {
|
||||
query: `SELECT table_name FROM information_schema.tables WHERE table_schema = ?`,
|
||||
params: [this.database_id],
|
||||
};
|
||||
}
|
||||
|
||||
getTableSchemaSql(table_name) {
|
||||
return {
|
||||
query: `SELECT COLUMN_NAME, COLUMN_TYPE, IS_NULLABLE, COLUMN_KEY, COLUMN_DEFAULT, EXTRA FROM information_schema.columns WHERE table_schema = ? AND table_name = ?`,
|
||||
params: [this.database_id, table_name],
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
module.exports.MySQLConnector = MySQLConnector;
|
||||
@@ -0,0 +1,75 @@
|
||||
const pgSql = require("pg");
|
||||
|
||||
class PostgresSQLConnector {
|
||||
#connected = false;
|
||||
constructor(
|
||||
config = {
|
||||
connectionString: null,
|
||||
schema: null,
|
||||
}
|
||||
) {
|
||||
this.className = "PostgresSQLConnector";
|
||||
this.connectionString = config.connectionString;
|
||||
this.schema = config.schema || "public";
|
||||
this._client = new pgSql.Client({
|
||||
connectionString: this.connectionString,
|
||||
});
|
||||
}
|
||||
|
||||
async connect() {
|
||||
await this._client.connect();
|
||||
this.#connected = true;
|
||||
return this._client;
|
||||
}
|
||||
|
||||
/**
|
||||
*
|
||||
* @param {string} queryString the SQL query to be run
|
||||
* @param {Array} params optional parameters for prepared statement
|
||||
* @returns {Promise<import(".").QueryResult>}
|
||||
*/
|
||||
async runQuery(queryString = "", params = []) {
|
||||
const result = { rows: [], count: 0, error: null };
|
||||
try {
|
||||
if (!this.#connected) await this.connect();
|
||||
const query = await this._client.query(queryString, params);
|
||||
result.rows = query.rows;
|
||||
result.count = query.rowCount;
|
||||
} catch (err) {
|
||||
console.log(this.className, err);
|
||||
result.error = err.message;
|
||||
} finally {
|
||||
// Check client is connected before closing since we use this for validation
|
||||
if (this._client) {
|
||||
await this._client.end();
|
||||
this.#connected = false;
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
async validateConnection() {
|
||||
try {
|
||||
const result = await this.runQuery("SELECT 1");
|
||||
return { success: !result.error, error: result.error };
|
||||
} catch (error) {
|
||||
return { success: false, error: error.message };
|
||||
}
|
||||
}
|
||||
|
||||
getTablesSql() {
|
||||
return {
|
||||
query: `SELECT * FROM pg_catalog.pg_tables WHERE schemaname = $1`,
|
||||
params: [this.schema],
|
||||
};
|
||||
}
|
||||
|
||||
getTableSchemaSql(table_name) {
|
||||
return {
|
||||
query: `SELECT column_name, data_type, character_maximum_length, column_default, is_nullable FROM INFORMATION_SCHEMA.COLUMNS WHERE table_name = $1 AND table_schema = $2`,
|
||||
params: [table_name, this.schema],
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
module.exports.PostgresSQLConnector = PostgresSQLConnector;
|
||||
@@ -0,0 +1,80 @@
|
||||
const { SystemSettings } = require("../../../../../../models/systemSettings");
|
||||
const { safeJsonParse } = require("../../../../../http");
|
||||
|
||||
/**
|
||||
* @typedef {('postgresql'|'mysql'|'sql-server')} SQLEngine
|
||||
*/
|
||||
|
||||
/**
|
||||
* @typedef {Object} QueryResult
|
||||
* @property {[number]} rows - The query result rows
|
||||
* @property {number} count - Number of rows the query returned/changed
|
||||
* @property {string|null} error - Error string if there was an issue
|
||||
*/
|
||||
|
||||
/**
|
||||
* A valid database SQL connection object
|
||||
* @typedef {Object} SQLConnection
|
||||
* @property {string} database_id - Unique identifier of the database connection
|
||||
* @property {SQLEngine} engine - Engine used by connection
|
||||
* @property {string} connectionString - RFC connection string for db
|
||||
*/
|
||||
|
||||
/**
|
||||
* @param {SQLEngine} identifier
|
||||
* @param {object} connectionConfig
|
||||
* @returns Database Connection Engine Class for SQLAgent or throws error
|
||||
*/
|
||||
function getDBClient(identifier = "", connectionConfig = {}) {
|
||||
switch (identifier) {
|
||||
case "mysql":
|
||||
const { MySQLConnector } = require("./MySQL");
|
||||
return new MySQLConnector(connectionConfig);
|
||||
case "postgresql":
|
||||
const { PostgresSQLConnector } = require("./Postgresql");
|
||||
return new PostgresSQLConnector(connectionConfig);
|
||||
case "sql-server":
|
||||
const { MSSQLConnector } = require("./MSSQL");
|
||||
return new MSSQLConnector(connectionConfig);
|
||||
default:
|
||||
throw new Error(
|
||||
`There is no supported database connector for ${identifier}`
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Lists all of the known database connection that can be used by the agent.
|
||||
* @returns {Promise<[SQLConnection]>}
|
||||
*/
|
||||
async function listSQLConnections() {
|
||||
return safeJsonParse(
|
||||
(await SystemSettings.get({ label: "agent_sql_connections" }))?.value,
|
||||
[]
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* Validates a SQL connection by attempting to connect and run a simple query
|
||||
* @param {SQLEngine} identifier - The SQL engine type
|
||||
* @param {object} connectionConfig - The connection configuration
|
||||
* @returns {Promise<{success: boolean, error: string|null}>}
|
||||
*/
|
||||
async function validateConnection(identifier = "", connectionConfig = {}) {
|
||||
try {
|
||||
const client = getDBClient(identifier, connectionConfig);
|
||||
return await client.validateConnection();
|
||||
} catch {
|
||||
console.log(`Failed to connect to ${identifier} database.`);
|
||||
return {
|
||||
success: false,
|
||||
error: `Unable to connect to ${identifier}. Please verify your connection details.`,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = {
|
||||
getDBClient,
|
||||
listSQLConnections,
|
||||
validateConnection,
|
||||
};
|
||||
@@ -0,0 +1,182 @@
|
||||
// Credit: https://github.com/sindilevich/connection-string-parser
|
||||
|
||||
/**
|
||||
* @typedef {Object} ConnectionStringParserOptions
|
||||
* @property {'mssql' | 'mysql' | 'postgresql' | 'db'} [scheme] - The scheme of the connection string
|
||||
*/
|
||||
|
||||
/**
|
||||
* @typedef {Object} ConnectionStringObject
|
||||
* @property {string} scheme - The scheme of the connection string eg: mongodb, mssql, mysql, postgresql, etc.
|
||||
* @property {string} username - The username of the connection string
|
||||
* @property {string} password - The password of the connection string
|
||||
* @property {{host: string, port: number}[]} hosts - The hosts of the connection string
|
||||
* @property {string} endpoint - The endpoint (database name) of the connection string
|
||||
* @property {Object} options - The options of the connection string
|
||||
*/
|
||||
class ConnectionStringParser {
|
||||
static DEFAULT_SCHEME = "db";
|
||||
|
||||
/**
|
||||
* @param {ConnectionStringParserOptions} options
|
||||
*/
|
||||
constructor(options = {}) {
|
||||
this.scheme =
|
||||
(options && options.scheme) || ConnectionStringParser.DEFAULT_SCHEME;
|
||||
}
|
||||
|
||||
/**
|
||||
* Takes a connection string object and returns a URI string of the form:
|
||||
*
|
||||
* scheme://[username[:password]@]host1[:port1][,host2[:port2],...[,hostN[:portN]]][/[endpoint]][?options]
|
||||
* @param {Object} connectionStringObject The object that describes connection string parameters
|
||||
*/
|
||||
format(connectionStringObject) {
|
||||
if (!connectionStringObject) {
|
||||
return this.scheme + "://localhost";
|
||||
}
|
||||
if (
|
||||
this.scheme &&
|
||||
connectionStringObject.scheme &&
|
||||
this.scheme !== connectionStringObject.scheme
|
||||
) {
|
||||
throw new Error(`Scheme not supported: ${connectionStringObject.scheme}`);
|
||||
}
|
||||
|
||||
let uri =
|
||||
(this.scheme ||
|
||||
connectionStringObject.scheme ||
|
||||
ConnectionStringParser.DEFAULT_SCHEME) + "://";
|
||||
|
||||
if (connectionStringObject.username) {
|
||||
uri += encodeURIComponent(connectionStringObject.username);
|
||||
// Allow empty passwords
|
||||
if (connectionStringObject.password) {
|
||||
uri += ":" + encodeURIComponent(connectionStringObject.password);
|
||||
}
|
||||
uri += "@";
|
||||
}
|
||||
uri += this._formatAddress(connectionStringObject);
|
||||
// Only put a slash when there is an endpoint
|
||||
if (connectionStringObject.endpoint) {
|
||||
uri += "/" + encodeURIComponent(connectionStringObject.endpoint);
|
||||
}
|
||||
if (
|
||||
connectionStringObject.options &&
|
||||
Object.keys(connectionStringObject.options).length > 0
|
||||
) {
|
||||
uri +=
|
||||
"?" +
|
||||
Object.keys(connectionStringObject.options)
|
||||
.map(
|
||||
(option) =>
|
||||
encodeURIComponent(option) +
|
||||
"=" +
|
||||
encodeURIComponent(connectionStringObject.options[option])
|
||||
)
|
||||
.join("&");
|
||||
}
|
||||
return uri;
|
||||
}
|
||||
|
||||
/**
|
||||
* Where scheme and hosts will always be present. Other fields will only be present in the result if they were
|
||||
* present in the input.
|
||||
* @param {string} uri The connection string URI
|
||||
* @returns {ConnectionStringObject} The connection string object
|
||||
*/
|
||||
parse(uri) {
|
||||
const connectionStringParser = new RegExp(
|
||||
"^\\s*" + // Optional whitespace padding at the beginning of the line
|
||||
"([^:]+)://" + // Scheme (Group 1)
|
||||
"(?:([^:@,/?=&]+)(?::([^:@,/?=&]+))?@)?" + // User (Group 2) and Password (Group 3)
|
||||
"([^@/?=&]+)" + // Host address(es) (Group 4)
|
||||
"(?:/([^:@,/?=&]+)?)?" + // Endpoint (Group 5)
|
||||
"(?:\\?([^:@,/?]+)?)?" + // Options (Group 6)
|
||||
"\\s*$", // Optional whitespace padding at the end of the line
|
||||
"gi"
|
||||
);
|
||||
const connectionStringObject = {};
|
||||
|
||||
if (!uri.includes("://")) {
|
||||
throw new Error(`No scheme found in URI ${uri}`);
|
||||
}
|
||||
|
||||
const tokens = connectionStringParser.exec(uri);
|
||||
|
||||
if (Array.isArray(tokens)) {
|
||||
connectionStringObject.scheme = tokens[1];
|
||||
if (this.scheme && this.scheme !== connectionStringObject.scheme) {
|
||||
throw new Error(`URI must start with '${this.scheme}://'`);
|
||||
}
|
||||
connectionStringObject.username = tokens[2]
|
||||
? decodeURIComponent(tokens[2])
|
||||
: tokens[2];
|
||||
connectionStringObject.password = tokens[3]
|
||||
? decodeURIComponent(tokens[3])
|
||||
: tokens[3];
|
||||
connectionStringObject.hosts = this._parseAddress(tokens[4]);
|
||||
connectionStringObject.endpoint = tokens[5]
|
||||
? decodeURIComponent(tokens[5])
|
||||
: tokens[5];
|
||||
connectionStringObject.options = tokens[6]
|
||||
? this._parseOptions(tokens[6])
|
||||
: tokens[6];
|
||||
}
|
||||
return connectionStringObject;
|
||||
}
|
||||
|
||||
/**
|
||||
* Formats the address portion of a connection string
|
||||
* @param {Object} connectionStringObject The object that describes connection string parameters
|
||||
*/
|
||||
_formatAddress(connectionStringObject) {
|
||||
return connectionStringObject.hosts
|
||||
.map(
|
||||
(address) =>
|
||||
encodeURIComponent(address.host) +
|
||||
(address.port
|
||||
? ":" + encodeURIComponent(address.port.toString(10))
|
||||
: "")
|
||||
)
|
||||
.join(",");
|
||||
}
|
||||
|
||||
/**
|
||||
* Parses an address
|
||||
* @param {string} addresses The address(es) to process
|
||||
*/
|
||||
_parseAddress(addresses) {
|
||||
return addresses.split(",").map((address) => {
|
||||
const i = address.indexOf(":");
|
||||
|
||||
return i >= 0
|
||||
? {
|
||||
host: decodeURIComponent(address.substring(0, i)),
|
||||
port: +address.substring(i + 1),
|
||||
}
|
||||
: { host: decodeURIComponent(address) };
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Parses options
|
||||
* @param {string} options The options to process
|
||||
*/
|
||||
_parseOptions(options) {
|
||||
const result = {};
|
||||
|
||||
options.split("&").forEach((option) => {
|
||||
const i = option.indexOf("=");
|
||||
|
||||
if (i >= 0) {
|
||||
result[decodeURIComponent(option.substring(0, i))] = decodeURIComponent(
|
||||
option.substring(i + 1)
|
||||
);
|
||||
}
|
||||
});
|
||||
return result;
|
||||
}
|
||||
}
|
||||
|
||||
module.exports = { ConnectionStringParser };
|
||||
@@ -0,0 +1,115 @@
|
||||
module.exports.SqlAgentGetTableSchema = {
|
||||
name: "sql-get-table-schema",
|
||||
plugin: function () {
|
||||
const {
|
||||
listSQLConnections,
|
||||
getDBClient,
|
||||
} = require("./SQLConnectors/index.js");
|
||||
|
||||
function formatQueryForDisplay(query, params = []) {
|
||||
if (!params.length) return query;
|
||||
let formatted = query;
|
||||
params.forEach((param, index) => {
|
||||
const value = typeof param === "string" ? `'${param}'` : param;
|
||||
formatted = formatted.replace(`$${index + 1}`, value);
|
||||
formatted = formatted.replace(`@p${index}`, value);
|
||||
formatted = formatted.replace("?", value);
|
||||
});
|
||||
return formatted;
|
||||
}
|
||||
|
||||
return {
|
||||
name: "sql-get-table-schema",
|
||||
setup(aibitat) {
|
||||
aibitat.function({
|
||||
super: aibitat,
|
||||
name: this.name,
|
||||
description:
|
||||
"Gets the table schema in SQL for a given `table` and `database_id`",
|
||||
examples: [
|
||||
{
|
||||
prompt: "What does the customers table in access-logs look like?",
|
||||
call: JSON.stringify({
|
||||
database_id: "access-logs",
|
||||
table_name: "customers",
|
||||
}),
|
||||
},
|
||||
{
|
||||
prompt:
|
||||
"Get me the full name of a company in records-main, the table should be call comps",
|
||||
call: JSON.stringify({
|
||||
database_id: "records-main",
|
||||
table_name: "comps",
|
||||
}),
|
||||
},
|
||||
],
|
||||
parameters: {
|
||||
$schema: "http://json-schema.org/draft-07/schema#",
|
||||
type: "object",
|
||||
properties: {
|
||||
database_id: {
|
||||
type: "string",
|
||||
description:
|
||||
"The database identifier for which we will connect to to query the table schema. This is a required field.",
|
||||
},
|
||||
table_name: {
|
||||
type: "string",
|
||||
description:
|
||||
"The database identifier for the table name we want the schema for. This is a required field.",
|
||||
},
|
||||
},
|
||||
additionalProperties: false,
|
||||
},
|
||||
required: ["database_id", "table_name"],
|
||||
handler: async function ({ database_id = "", table_name = "" }) {
|
||||
this.super.handlerProps.log(`Using the sql-get-table-schema tool.`);
|
||||
try {
|
||||
const databaseConfig = (await listSQLConnections()).find(
|
||||
(db) => db.database_id === database_id
|
||||
);
|
||||
if (!databaseConfig) {
|
||||
this.super.handlerProps.log(
|
||||
`sql-get-table-schema to find config!`,
|
||||
database_id
|
||||
);
|
||||
return `No database connection for ${database_id} was found!`;
|
||||
}
|
||||
|
||||
const db = getDBClient(databaseConfig.engine, databaseConfig);
|
||||
this.super.introspect(
|
||||
`${this.caller}: Querying the table schema for ${table_name} in the ${databaseConfig.database_id} database.`
|
||||
);
|
||||
|
||||
const sqlQuery = db.getTableSchemaSql(table_name);
|
||||
const isParameterized =
|
||||
typeof sqlQuery === "object" && sqlQuery.query;
|
||||
const queryString = isParameterized ? sqlQuery.query : sqlQuery;
|
||||
const queryParams = isParameterized ? sqlQuery.params : [];
|
||||
|
||||
this.super.introspect(
|
||||
`Running SQL: ${formatQueryForDisplay(queryString, queryParams)}`
|
||||
);
|
||||
const result = await db.runQuery(queryString, queryParams);
|
||||
|
||||
if (result.error) {
|
||||
this.super.handlerProps.log(
|
||||
`sql-get-table-schema tool reported error`,
|
||||
result.error
|
||||
);
|
||||
this.super.introspect(`Error: ${result.error}`);
|
||||
return `There was an error running the query: ${result.error}`;
|
||||
}
|
||||
|
||||
return JSON.stringify(result);
|
||||
} catch (e) {
|
||||
this.super.handlerProps.log(
|
||||
`sql-get-table-schema raised an error. ${e.message}`
|
||||
);
|
||||
return e.message;
|
||||
}
|
||||
},
|
||||
});
|
||||
},
|
||||
};
|
||||
},
|
||||
};
|
||||
@@ -0,0 +1,21 @@
|
||||
const { SqlAgentGetTableSchema } = require("./get-table-schema");
|
||||
const { SqlAgentListDatabase } = require("./list-database");
|
||||
const { SqlAgentListTables } = require("./list-table");
|
||||
const { SqlAgentQuery } = require("./query");
|
||||
|
||||
const sqlAgent = {
|
||||
name: "sql-agent",
|
||||
startupConfig: {
|
||||
params: {},
|
||||
},
|
||||
plugin: [
|
||||
SqlAgentListDatabase,
|
||||
SqlAgentListTables,
|
||||
SqlAgentGetTableSchema,
|
||||
SqlAgentQuery,
|
||||
],
|
||||
};
|
||||
|
||||
module.exports = {
|
||||
sqlAgent,
|
||||
};
|
||||
@@ -0,0 +1,49 @@
|
||||
module.exports.SqlAgentListDatabase = {
|
||||
name: "sql-list-databases",
|
||||
plugin: function () {
|
||||
const { listSQLConnections } = require("./SQLConnectors");
|
||||
return {
|
||||
name: "sql-list-databases",
|
||||
setup(aibitat) {
|
||||
aibitat.function({
|
||||
super: aibitat,
|
||||
name: this.name,
|
||||
description:
|
||||
"List all available databases via `list_databases` you currently have access to. Returns a unique string identifier `database_id` that can be used for future calls.",
|
||||
examples: [
|
||||
{
|
||||
prompt: "What databases can you access?",
|
||||
call: JSON.stringify({}),
|
||||
},
|
||||
{
|
||||
prompt: "What databases can you tell me about?",
|
||||
call: JSON.stringify({}),
|
||||
},
|
||||
{
|
||||
prompt: "Is there a database named erp-logs you can access?",
|
||||
call: JSON.stringify({}),
|
||||
},
|
||||
],
|
||||
parameters: {
|
||||
$schema: "http://json-schema.org/draft-07/schema#",
|
||||
type: "object",
|
||||
properties: {},
|
||||
additionalProperties: false,
|
||||
},
|
||||
handler: async function () {
|
||||
this.super.handlerProps.log(`Using the sql-list-databases tool.`);
|
||||
this.super.introspect(
|
||||
`${this.caller}: Checking what are the available databases.`
|
||||
);
|
||||
|
||||
const connections = (await listSQLConnections()).map((conn) => {
|
||||
const { connectionString: _connectionString, ...rest } = conn;
|
||||
return rest;
|
||||
});
|
||||
return JSON.stringify(connections);
|
||||
},
|
||||
});
|
||||
},
|
||||
};
|
||||
},
|
||||
};
|
||||
@@ -0,0 +1,105 @@
|
||||
module.exports.SqlAgentListTables = {
|
||||
name: "sql-list-tables",
|
||||
plugin: function () {
|
||||
const {
|
||||
listSQLConnections,
|
||||
getDBClient,
|
||||
} = require("./SQLConnectors/index.js");
|
||||
|
||||
function formatQueryForDisplay(query, params = []) {
|
||||
if (!params.length) return query;
|
||||
let formatted = query;
|
||||
params.forEach((param, index) => {
|
||||
const value = typeof param === "string" ? `'${param}'` : param;
|
||||
formatted = formatted.replace(`$${index + 1}`, value);
|
||||
formatted = formatted.replace(`@p${index}`, value);
|
||||
formatted = formatted.replace("?", value);
|
||||
});
|
||||
return formatted;
|
||||
}
|
||||
|
||||
return {
|
||||
name: "sql-list-tables",
|
||||
setup(aibitat) {
|
||||
aibitat.function({
|
||||
super: aibitat,
|
||||
name: this.name,
|
||||
description:
|
||||
"List all available tables in a database via its `database_id`.",
|
||||
examples: [
|
||||
{
|
||||
prompt: "What tables are there in the `access-logs` database?",
|
||||
call: JSON.stringify({ database_id: "access-logs" }),
|
||||
},
|
||||
{
|
||||
prompt:
|
||||
"What information can you access in the customer_accts postgres db?",
|
||||
call: JSON.stringify({ database_id: "customer_accts" }),
|
||||
},
|
||||
{
|
||||
prompt: "Can you tell me what is in the primary-logs db?",
|
||||
call: JSON.stringify({ database_id: "primary-logs" }),
|
||||
},
|
||||
],
|
||||
parameters: {
|
||||
$schema: "http://json-schema.org/draft-07/schema#",
|
||||
type: "object",
|
||||
properties: {
|
||||
database_id: {
|
||||
type: "string",
|
||||
description:
|
||||
"The database identifier for which we will list all tables for. This is a required parameter",
|
||||
},
|
||||
},
|
||||
additionalProperties: false,
|
||||
},
|
||||
required: ["database_id"],
|
||||
handler: async function ({ database_id = "" }) {
|
||||
try {
|
||||
this.super.handlerProps.log(`Using the sql-list-tables tool.`);
|
||||
const databaseConfig = (await listSQLConnections()).find(
|
||||
(db) => db.database_id === database_id
|
||||
);
|
||||
if (!databaseConfig) {
|
||||
this.super.handlerProps.log(
|
||||
`sql-list-tables failed to find config!`,
|
||||
database_id
|
||||
);
|
||||
return `No database connection for ${database_id} was found!`;
|
||||
}
|
||||
|
||||
const db = getDBClient(databaseConfig.engine, databaseConfig);
|
||||
this.super.introspect(
|
||||
`${this.caller}: Checking what are the available tables in the ${databaseConfig.database_id} database.`
|
||||
);
|
||||
|
||||
const sqlQuery = db.getTablesSql();
|
||||
const isParameterized =
|
||||
typeof sqlQuery === "object" && sqlQuery.query;
|
||||
const queryString = isParameterized ? sqlQuery.query : sqlQuery;
|
||||
const queryParams = isParameterized ? sqlQuery.params : [];
|
||||
|
||||
this.super.introspect(
|
||||
`Running SQL: ${formatQueryForDisplay(queryString, queryParams)}`
|
||||
);
|
||||
const result = await db.runQuery(queryString, queryParams);
|
||||
if (result.error) {
|
||||
this.super.handlerProps.log(
|
||||
`sql-list-tables tool reported error`,
|
||||
result.error
|
||||
);
|
||||
this.super.introspect(`Error: ${result.error}`);
|
||||
return `There was an error running the query: ${result.error}`;
|
||||
}
|
||||
|
||||
return JSON.stringify(result);
|
||||
} catch (e) {
|
||||
console.error(e);
|
||||
return e.message;
|
||||
}
|
||||
},
|
||||
});
|
||||
},
|
||||
};
|
||||
},
|
||||
};
|
||||
@@ -0,0 +1,101 @@
|
||||
module.exports.SqlAgentQuery = {
|
||||
name: "sql-query",
|
||||
plugin: function () {
|
||||
const {
|
||||
getDBClient,
|
||||
listSQLConnections,
|
||||
} = require("./SQLConnectors/index.js");
|
||||
|
||||
return {
|
||||
name: "sql-query",
|
||||
setup(aibitat) {
|
||||
aibitat.function({
|
||||
super: aibitat,
|
||||
name: this.name,
|
||||
description:
|
||||
"Run a read-only SQL query on a `database_id` which will return up rows of data related to the query. The query must only be SELECT statements which do not modify the table data. There should be a reasonable LIMIT on the return quantity to prevent long-running or queries which crash the db.",
|
||||
examples: [
|
||||
{
|
||||
prompt: "How many customers are in dvd-rentals?",
|
||||
call: JSON.stringify({
|
||||
database_id: "dvd-rentals",
|
||||
sql_query: "SELECT * FROM customers",
|
||||
}),
|
||||
},
|
||||
{
|
||||
prompt: "Can you tell me the total volume of sales last month?",
|
||||
call: JSON.stringify({
|
||||
database_id: "sales-db",
|
||||
sql_query:
|
||||
"SELECT SUM(sale_amount) AS total_sales FROM sales WHERE sale_date >= DATEADD(month, -1, DATEFROMPARTS(YEAR(GETDATE()), MONTH(GETDATE()), 1)) AND sale_date < DATEFROMPARTS(YEAR(GETDATE()), MONTH(GETDATE()), 1)",
|
||||
}),
|
||||
},
|
||||
{
|
||||
prompt:
|
||||
"Do we have anyone in the staff table for our production db named 'sam'? ",
|
||||
call: JSON.stringify({
|
||||
database_id: "production",
|
||||
sql_query:
|
||||
"SElECT * FROM staff WHERE first_name='sam%' OR last_name='sam%'",
|
||||
}),
|
||||
},
|
||||
],
|
||||
parameters: {
|
||||
$schema: "http://json-schema.org/draft-07/schema#",
|
||||
type: "object",
|
||||
properties: {
|
||||
database_id: {
|
||||
type: "string",
|
||||
description:
|
||||
"The database identifier for which we will connect to to query the table schema. This is required to run the SQL query.",
|
||||
},
|
||||
sql_query: {
|
||||
type: "string",
|
||||
description:
|
||||
"The raw SQL query to run. Should be a query which does not modify the table and will return results.",
|
||||
},
|
||||
},
|
||||
additionalProperties: false,
|
||||
},
|
||||
required: ["database_id", "table_name"],
|
||||
handler: async function ({ database_id = "", sql_query = "" }) {
|
||||
this.super.handlerProps.log(`Using the sql-query tool.`);
|
||||
try {
|
||||
const databaseConfig = (await listSQLConnections()).find(
|
||||
(db) => db.database_id === database_id
|
||||
);
|
||||
if (!databaseConfig) {
|
||||
this.super.handlerProps.log(
|
||||
`sql-query failed to find config!`,
|
||||
database_id
|
||||
);
|
||||
return `No database connection for ${database_id} was found!`;
|
||||
}
|
||||
|
||||
this.super.introspect(
|
||||
`${this.caller}: Im going to run a query on the ${database_id} to get an answer.`
|
||||
);
|
||||
const db = getDBClient(databaseConfig.engine, databaseConfig);
|
||||
|
||||
this.super.introspect(`Running SQL: ${sql_query}`);
|
||||
const result = await db.runQuery(sql_query);
|
||||
if (result.error) {
|
||||
this.super.handlerProps.log(
|
||||
`sql-query tool reported error`,
|
||||
result.error
|
||||
);
|
||||
this.super.introspect(`Error: ${result.error}`);
|
||||
return `There was an error running the query: ${result.error}`;
|
||||
}
|
||||
|
||||
return JSON.stringify(result);
|
||||
} catch (e) {
|
||||
console.error(e);
|
||||
return e.message;
|
||||
}
|
||||
},
|
||||
});
|
||||
},
|
||||
};
|
||||
},
|
||||
};
|
||||
Reference in New Issue
Block a user