@godspeedsystems/plugins-text-to-sql-as-datasource
Version:
text-to-sql as datasource plugin for Godspeed Framework
563 lines (561 loc) • 23.3 kB
JavaScript
"use strict";
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
var desc = Object.getOwnPropertyDescriptor(m, k);
if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
desc = { enumerable: true, get: function() { return m[k]; } };
}
Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
o["default"] = v;
});
var __importStar = (this && this.__importStar) || function (mod) {
if (mod && mod.__esModule) return mod;
var result = {};
if (mod != null) for (var k in mod) if (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
__setModuleDefault(result, mod);
return result;
};
var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
return new (P || (P = Promise))(function (resolve, reject) {
function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
step((generator = generator.apply(thisArg, _arguments || [])).next());
});
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.DEFAULT_CONFIG = exports.CONFIG_FILE_NAME = exports.Type = exports.SourceType = exports.DataSource = void 0;
const core_1 = require("@godspeedsystems/core");
const pg_1 = require("pg");
const promise_1 = require("mysql2/promise");
const mongodb_1 = require("mongodb");
const oracledb = __importStar(require("oracledb"));
const redis_1 = require("redis");
const crypto_1 = require("crypto");
const generative_ai_1 = require("@google/generative-ai");
class DatabaseConnectionManager {
constructor() {
this.connections = new Map();
}
connect(config) {
return __awaiter(this, void 0, void 0, function* () {
try {
switch (config.type) {
case 'postgres':
return yield this.connectPostgres(config.config);
case 'mysql':
return yield this.connectMySQL(config.config);
case 'mongodb':
return yield this.connectMongo(config.config);
case 'oracle':
return yield this.connectOracle(config.config);
default:
throw new Error(`Database not supported yet: ${config.type}`);
}
}
catch (error) {
core_1.logger.error(`Database connection failed for ${config.type}:`, error);
throw error;
}
});
}
connectPostgres(config) {
return __awaiter(this, void 0, void 0, function* () {
try {
const pool = new pg_1.Pool(config);
yield pool.query('SELECT 1');
core_1.logger.info('PostgreSQL connection established');
return { client: pool, type: 'postgres', isConnected: true };
}
catch (error) {
core_1.logger.error('PostgreSQL connection failed:', error);
throw error;
}
});
}
connectMySQL(config) {
return __awaiter(this, void 0, void 0, function* () {
try {
const connection = yield (0, promise_1.createConnection)(config);
yield connection.query('SELECT 1');
core_1.logger.info('MySQL connection established');
return { client: connection, type: 'mysql', isConnected: true };
}
catch (error) {
core_1.logger.error('MySQL connection failed:', error);
throw error;
}
});
}
connectMongo(config) {
return __awaiter(this, void 0, void 0, function* () {
try {
const client = yield mongodb_1.MongoClient.connect(config.url);
const db = client.db(config.database);
core_1.logger.info('MongoDB connection established');
return { client: db, type: 'mongodb', isConnected: true };
}
catch (error) {
core_1.logger.error('MongoDB connection failed:', error);
throw error;
}
});
}
connectOracle(config) {
return __awaiter(this, void 0, void 0, function* () {
try {
const connection = yield oracledb.getConnection({
user: config.user,
password: config.password,
connectString: config.connectString,
});
yield connection.execute('SELECT 1 FROM DUAL');
core_1.logger.info('Oracle connection established');
return { client: connection, type: 'oracle', isConnected: true };
}
catch (error) {
core_1.logger.error('Oracle connection failed:', error);
throw error;
}
});
}
}
class MultiDBTextToSQLDataSource extends core_1.GSDataSource {
initClient() {
return __awaiter(this, void 0, void 0, function* () {
try {
core_1.logger.info('Initializing Text-to-SQL service');
this.genAI = new generative_ai_1.GoogleGenerativeAI(process.env.GEMINI_API_KEY);
this.model = this.genAI.getGenerativeModel({ model: 'gemini-pro' });
this.dbManager = new DatabaseConnectionManager();
this.connections = new Map();
this.schemas = new Map();
yield this.initializeDatabases();
this.redisClient = (0, redis_1.createClient)({
url: process.env.REDIS_URL,
});
core_1.logger.info('Text-to-SQL service initialized successfully');
return { status: 'initialized' };
}
catch (error) {
core_1.logger.error('Service initialization failed:', error);
throw error;
}
});
}
initializeDatabases() {
return __awaiter(this, void 0, void 0, function* () {
const dbConfigs = [
{
type: 'postgres',
config: {
user: process.env.PG_USER,
host: process.env.PG_HOST,
database: process.env.PG_DB,
password: process.env.PG_PASSWORD,
port: parseInt(process.env.PG_PORT || '5432'),
},
},
{
type: 'mysql',
config: {
host: process.env.MYSQL_HOST,
user: process.env.MYSQL_USER,
password: process.env.MYSQL_PASSWORD,
database: process.env.MYSQL_DB,
},
},
{
type: 'mongodb',
config: {
url: process.env.MONGODB_URL,
database: process.env.MONGODB_DB,
},
},
{
type: 'oracle',
config: {
user: process.env.ORACLE_USER,
password: process.env.ORACLE_PASSWORD,
connectString: process.env.ORACLE_CONNECT_STRING,
},
},
];
for (const config of dbConfigs) {
try {
core_1.logger.info(`Initializing ${config.type} connection`);
const connection = yield this.dbManager.connect(config);
this.connections.set(config.type, connection);
this.schemas.set(config.type, yield this.fetchDatabaseSchema(config.type));
core_1.logger.info(`${config.type} initialization complete`);
}
catch (error) {
core_1.logger.error(`Failed to initialize ${config.type}:`, error);
}
}
});
}
fetchDatabaseSchema(dbType) {
return __awaiter(this, void 0, void 0, function* () {
const connection = this.connections.get(dbType);
if (!connection)
throw new Error(`No connection found for ${dbType}`);
try {
core_1.logger.info(`Fetching schema for ${dbType}`);
switch (dbType) {
case 'postgres':
return yield this.fetchPostgresSchema(connection);
case 'mysql':
return yield this.fetchMySQLSchema(connection);
case 'mongodb':
return yield this.fetchMongoSchema(connection);
case 'oracle':
return yield this.fetchOracleSchema(connection);
default:
throw new Error(`Unsupported database type: ${dbType}`);
}
}
catch (error) {
core_1.logger.error(`Schema fetch failed for ${dbType}:`, error);
throw error;
}
});
}
fetchPostgresSchema(connection) {
return __awaiter(this, void 0, void 0, function* () {
const schemaQuery = `
SELECT
table_name,
column_name,
data_type,
is_nullable
FROM
information_schema.columns
WHERE
table_schema = 'public'
ORDER BY
table_name, ordinal_position;
`;
const result = yield connection.client.query(schemaQuery);
return this.formatRelationalSchema(result.rows);
});
}
fetchMySQLSchema(connection) {
return __awaiter(this, void 0, void 0, function* () {
const [rows] = yield connection.client.query(`
SELECT
TABLE_NAME,
COLUMN_NAME,
DATA_TYPE,
IS_NULLABLE
FROM
INFORMATION_SCHEMA.COLUMNS
WHERE
TABLE_SCHEMA = DATABASE()
`);
return this.formatRelationalSchema(rows);
});
}
fetchMongoSchema(connection) {
return __awaiter(this, void 0, void 0, function* () {
const db = connection.client;
const collections = yield db.listCollections().toArray();
return this.formatMongoSchema(collections);
});
}
fetchOracleSchema(connection) {
return __awaiter(this, void 0, void 0, function* () {
const result = yield connection.client.execute(`
SELECT
table_name,
column_name,
data_type,
nullable,
data_length,
data_precision
FROM
user_tab_columns
ORDER BY
table_name, column_id
`, [], { outFormat: oracledb.OUT_FORMAT_OBJECT });
return this.formatOracleSchema(result.rows);
});
}
formatRelationalSchema(rows) {
const tableMap = new Map();
rows.forEach((row) => {
var _a;
const tableName = row.table_name || row.TABLE_NAME;
if (!tableMap.has(tableName)) {
tableMap.set(tableName, []);
}
(_a = tableMap.get(tableName)) === null || _a === void 0 ? void 0 : _a.push(row);
});
return Array.from(tableMap.entries())
.map(([table, columns]) => {
return `Table ${table} {\n${columns
.map((col) => ` ${col.column_name || col.COLUMN_NAME} ${col.data_type || col.DATA_TYPE}`)
.join('\n')}\n}`;
})
.join('\n\n');
}
formatMongoSchema(collections) {
return collections
.map((collection) => `Collection ${collection.name} {\n // Schema is dynamic\n}`)
.join('\n\n');
}
formatOracleSchema(rows) {
const tableMap = new Map();
rows.forEach((row) => {
var _a;
if (!tableMap.has(row.TABLE_NAME)) {
tableMap.set(row.TABLE_NAME, []);
}
(_a = tableMap.get(row.TABLE_NAME)) === null || _a === void 0 ? void 0 : _a.push(row);
});
return Array.from(tableMap.entries())
.map(([table, columns]) => {
return `Table ${table} {\n${columns
.map((col) => ` ${col.COLUMN_NAME} ${col.DATA_TYPE}(${col.DATA_LENGTH})`)
.join('\n')}\n}`;
})
.join('\n\n');
}
execute(ctx, args) {
return __awaiter(this, void 0, void 0, function* () {
const { query, dbType = 'postgres', //defaults to postgres
validateOnly = false, cache = true, } = args;
try {
core_1.logger.info('Executing query', { dbType, validateOnly });
const connection = this.connections.get(dbType);
if (!connection) {
throw new Error(`No connection available for database type: ${dbType}`);
}
if (cache) {
const cachedResult = yield this.getCachedQuery(query, dbType);
if (cachedResult) {
core_1.logger.info('Cache hit, returning cached result');
return cachedResult;
}
}
const generatedQuery = yield this.generateQuery(query, dbType);
core_1.logger.debug('Generated SQL query', { generatedQuery });
if (validateOnly) {
return { status: 'valid', query: generatedQuery };
}
const result = yield this.executeQuery(generatedQuery, dbType);
if (cache) {
yield this.cacheQuery(query, dbType, result);
}
core_1.logger.info('Query executed successfully');
return result;
}
catch (error) {
core_1.logger.error('Query execution failed:', error);
throw error;
}
});
}
generateQuery(naturalQuery, dbType) {
return __awaiter(this, void 0, void 0, function* () {
const schema = this.schemas.get(dbType);
try {
const prompt = `
Schema:
${schema}
Task: Convert this natural language query to a PostgreSQL query:
"${naturalQuery}"
Requirements:
- Return only the raw SQL query
- No markdown formatting
- No explanations
- Must start with SELECT
- Only use existing tables and columns
`;
const result = yield this.model.generateContent(prompt);
let sql = result.response.text().trim();
// Clean up the response
sql = sql.replace(/```sql/gi, '')
.replace(/```/g, '')
.replace(/`/g, '')
.trim();
// Basic validation
if (!sql.toLowerCase().startsWith('select')) {
throw new Error('Generated query must start with SELECT');
}
return sql;
}
catch (error) {
core_1.logger.error(`SQL generation failed: ${error.message}`);
throw new Error(`Failed to generate SQL: ${error.message}`);
}
});
}
executeQuery(query, dbType) {
return __awaiter(this, void 0, void 0, function* () {
const connection = this.connections.get(dbType);
if (!connection)
throw new Error(`No connection found for ${dbType}`);
try {
switch (dbType) {
case 'postgres':
const pgResult = yield connection.client.query(query);
return {
data: pgResult.rows,
metadata: { rowCount: pgResult.rowCount },
};
case 'mysql':
const [mysqlRows] = yield connection.client.query(query);
return { data: mysqlRows };
case 'mongodb':
const mongoResult = yield eval(`connection.client.${query}`);
return {
data: Array.isArray(mongoResult) ? mongoResult : [mongoResult],
};
case 'oracle':
const oracleResult = yield connection.client.execute(query, [], {
outFormat: oracledb.OUT_FORMAT_OBJECT,
});
return {
data: oracleResult.rows || [],
metadata: {
rowsAffected: oracleResult.rowsAffected,
metaData: oracleResult.metaData,
},
};
default:
throw new Error(`Unsupported database type: ${dbType}`);
}
}
catch (error) {
core_1.logger.error(`Query execution failed for ${dbType}:`, error);
throw error;
}
});
}
cacheQuery(query, dbType, result) {
return __awaiter(this, void 0, void 0, function* () {
const hash = (0, crypto_1.createHash)('md5').update(query).update(dbType).digest('hex');
try {
yield this.redisClient.set(`sql_cache:${hash}`, JSON.stringify(result), {
EX: 3600, // Sets expiration to 1 hour
});
core_1.logger.debug('Query result cached', { hash });
}
catch (error) {
core_1.logger.warn('Cache operation failed:', error);
}
});
}
getCachedQuery(query, dbType) {
return __awaiter(this, void 0, void 0, function* () {
const hash = (0, crypto_1.createHash)('md5').update(query).update(dbType).digest('hex');
try {
const cached = yield this.redisClient.get(`sql_cache:${hash}`);
return cached ? JSON.parse(cached) : null;
}
catch (error) {
core_1.logger.warn('Cache retrieval failed:', error);
return null;
}
});
}
cleanup() {
return __awaiter(this, void 0, void 0, function* () {
try {
core_1.logger.info('Starting cleanup');
// Close other database connections
for (const [dbType, connection] of this.connections.entries()) {
try {
switch (connection.type) {
case 'postgres':
yield connection.client.end();
break;
case 'mysql':
yield connection.client.end();
break;
case 'oracle':
yield connection.client.close();
break;
case 'mongodb':
yield connection.client.close();
break;
}
core_1.logger.info(`Closed ${dbType} connection`);
}
catch (error) {
core_1.logger.error(`Failed to close ${dbType} connection:`, error);
}
}
yield this.redisClient.quit(); // Ensure Redis client is properly closed
core_1.logger.info('Cleanup completed successfully');
}
catch (error) {
core_1.logger.error('Cleanup failed:', error);
throw error;
}
});
}
}
exports.default = MultiDBTextToSQLDataSource;
exports.DataSource = MultiDBTextToSQLDataSource;
const SourceType = 'DS';
exports.SourceType = SourceType;
const Type = 'text-to-sql';
exports.Type = Type;
const CONFIG_FILE_NAME = 'text-to-sql';
exports.CONFIG_FILE_NAME = CONFIG_FILE_NAME;
const DEFAULT_CONFIG = {
gemini: {
apiKey: process.env.GEMINI_API_KEY,
},
databases: {
postgres: {
enabled: true,
config: {
user: process.env.PG_USER,
host: process.env.PG_HOST,
database: process.env.PG_DB,
password: process.env.PG_PASSWORD,
port: parseInt(process.env.PG_PORT || '5432'),
},
},
mysql: {
enabled: true,
config: {
host: process.env.MYSQL_HOST,
user: process.env.MYSQL_USER,
password: process.env.MYSQL_PASSWORD,
database: process.env.MYSQL_DB,
},
},
mongodb: {
enabled: true,
config: {
url: process.env.MONGODB_URL,
database: process.env.MONGODB_DB,
},
},
oracle: {
enabled: true,
config: {
user: process.env.ORACLE_USER,
password: process.env.ORACLE_PASSWORD,
connectString: process.env.ORACLE_CONNECT_STRING,
},
},
},
redis: {
url: process.env.REDIS_URL,
},
};
exports.DEFAULT_CONFIG = DEFAULT_CONFIG;
//# sourceMappingURL=index.js.map