mirror of https://github.com/sysown/proxysql
mcp discovery: enforce target-scoped run model and add protocol-aware static harvesting (mysql+pgsql)
Refactor the MCP query/discovery stack to remove prototype-era MySQL-only assumptions and make discovery/catalog semantics explicitly target-scoped. This commit is intentionally not backward compatible with legacy single-target catalog metadata. Key implementation details: - Discovery_Schema: - runs table now stores target_id, protocol, and server_version - resolve_run_id() now requires target_id and resolves schemas within target scope - create_run() now records target/protocol/server version - added legacy schema detection + destructive catalog table rebuild when old runs layout is found - fixed resultset lifetime/escaping issues in target-aware run resolution - New PostgreSQL static harvester: - added PgSQL_Static_Harvester class and build wiring - harvests schemas/objects/columns/indexes/fks/view definitions into Discovery_Schema - maps pg metadata to object_id-based insert APIs used by Discovery_Schema - Query_Tool_Handler: - constructor simplified to catalog_path only (no mcp-mysql_* runtime ctor deps) - discovery.run_static now requires target_id and dispatches protocol-aware harvester at runtime - catalog.*, agent.run_start, and llm.* tool contracts now require target_id when resolving run_id - all run_id resolution call sites switched to resolve_run_id(target_id, ...) - list_schemas catalog query now scoped by target_id via runs join - digest logging path now resolves run_id in target context - ProxySQL_MCP_Server: - updated Query_Tool_Handler construction to new signature - Docs updated: - added explicit target_id scoping guidance for discovery/catalog/agent/llm workflows - clarified protocol-aware routing and noted legacy examples still present in large script docs Build validation: - Recompiled changed objects with PROXYSQLGENAI=1: - Query_Tool_Handler.oo - Discovery_Schema.oo - PgSQL_Static_Harvester.oo - ProxySQL_MCP_Server.oov3.0-MCP_multi
Rene Cannao
5 days ago
parent
7e54883ba5
commit
9685cdaa4b
@ -0,0 +1,70 @@
|
||||
#ifndef CLASS_PGSQL_STATIC_HARVESTER_H
|
||||
#define CLASS_PGSQL_STATIC_HARVESTER_H
|
||||
|
||||
#ifdef PROXYSQLGENAI
|
||||
|
||||
#include "Discovery_Schema.h"
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <memory>
|
||||
#include <pthread.h>
|
||||
|
||||
typedef struct pg_conn PGconn;
|
||||
|
||||
class PgSQL_Static_Harvester {
|
||||
private:
|
||||
std::string pgsql_host;
|
||||
int pgsql_port;
|
||||
std::string pgsql_user;
|
||||
std::string pgsql_password;
|
||||
std::string pgsql_dbname;
|
||||
PGconn* pgsql_conn;
|
||||
pthread_mutex_t conn_lock;
|
||||
|
||||
Discovery_Schema* catalog;
|
||||
|
||||
int current_run_id;
|
||||
std::string source_dsn;
|
||||
std::string server_version;
|
||||
|
||||
int connect_pgsql();
|
||||
void disconnect_pgsql();
|
||||
int execute_query(const std::string& query, std::vector<std::vector<std::string>>& results);
|
||||
std::string get_pgsql_version();
|
||||
static bool is_valid_schema_name(const std::string& name);
|
||||
static std::string escape_sql_string(const std::string& str);
|
||||
|
||||
public:
|
||||
PgSQL_Static_Harvester(
|
||||
const std::string& host,
|
||||
int port,
|
||||
const std::string& user,
|
||||
const std::string& password,
|
||||
const std::string& dbname,
|
||||
const std::string& catalog_path
|
||||
);
|
||||
~PgSQL_Static_Harvester();
|
||||
|
||||
int init();
|
||||
void close();
|
||||
|
||||
int start_run(const std::string& target_id, const std::string& notes = "");
|
||||
int finish_run(const std::string& notes = "");
|
||||
int get_run_id() const { return current_run_id; }
|
||||
|
||||
int harvest_schemas(const std::string& only_schema = "");
|
||||
int harvest_objects(const std::string& only_schema = "");
|
||||
int harvest_columns(const std::string& only_schema = "");
|
||||
int harvest_indexes(const std::string& only_schema = "");
|
||||
int harvest_foreign_keys(const std::string& only_schema = "");
|
||||
int harvest_view_definitions(const std::string& only_schema = "");
|
||||
int build_quick_profiles();
|
||||
int rebuild_fts_index();
|
||||
|
||||
int run_full_harvest(const std::string& target_id, const std::string& only_schema = "", const std::string& notes = "");
|
||||
std::string get_harvest_stats();
|
||||
std::string get_harvest_stats(int run_id);
|
||||
};
|
||||
|
||||
#endif /* PROXYSQLGENAI */
|
||||
#endif /* CLASS_PGSQL_STATIC_HARVESTER_H */
|
||||
@ -0,0 +1,699 @@
|
||||
#ifdef PROXYSQLGENAI
|
||||
|
||||
#include "PgSQL_Static_Harvester.h"
|
||||
#include "proxysql.h"
|
||||
#include "cpp.h"
|
||||
|
||||
#include <libpq-fe.h>
|
||||
#include <sstream>
|
||||
#include <map>
|
||||
#include <algorithm>
|
||||
#include <cctype>
|
||||
#include "../deps/json/json.hpp"
|
||||
|
||||
using json = nlohmann::json;
|
||||
|
||||
PgSQL_Static_Harvester::PgSQL_Static_Harvester(
|
||||
const std::string& host,
|
||||
int port,
|
||||
const std::string& user,
|
||||
const std::string& password,
|
||||
const std::string& dbname,
|
||||
const std::string& catalog_path
|
||||
) : pgsql_host(host),
|
||||
pgsql_port(port),
|
||||
pgsql_user(user),
|
||||
pgsql_password(password),
|
||||
pgsql_dbname(dbname),
|
||||
pgsql_conn(NULL),
|
||||
catalog(NULL),
|
||||
current_run_id(-1)
|
||||
{
|
||||
pthread_mutex_init(&conn_lock, NULL);
|
||||
catalog = new Discovery_Schema(catalog_path);
|
||||
}
|
||||
|
||||
PgSQL_Static_Harvester::~PgSQL_Static_Harvester() {
|
||||
close();
|
||||
if (catalog) {
|
||||
delete catalog;
|
||||
catalog = NULL;
|
||||
}
|
||||
pthread_mutex_destroy(&conn_lock);
|
||||
}
|
||||
|
||||
int PgSQL_Static_Harvester::init() {
|
||||
if (catalog->init()) {
|
||||
proxy_error("PgSQL_Static_Harvester: Failed to initialize catalog\n");
|
||||
return -1;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
void PgSQL_Static_Harvester::close() {
|
||||
disconnect_pgsql();
|
||||
}
|
||||
|
||||
int PgSQL_Static_Harvester::connect_pgsql() {
|
||||
pthread_mutex_lock(&conn_lock);
|
||||
|
||||
if (pgsql_conn) {
|
||||
pthread_mutex_unlock(&conn_lock);
|
||||
return 0;
|
||||
}
|
||||
|
||||
std::ostringstream conninfo;
|
||||
conninfo << "host=" << pgsql_host
|
||||
<< " port=" << pgsql_port
|
||||
<< " user=" << pgsql_user
|
||||
<< " password=" << pgsql_password
|
||||
<< " connect_timeout=30";
|
||||
if (!pgsql_dbname.empty()) {
|
||||
conninfo << " dbname=" << pgsql_dbname;
|
||||
}
|
||||
|
||||
pgsql_conn = PQconnectdb(conninfo.str().c_str());
|
||||
if (pgsql_conn == NULL || PQstatus(pgsql_conn) != CONNECTION_OK) {
|
||||
proxy_error("PgSQL_Static_Harvester: PQconnectdb failed: %s\n",
|
||||
pgsql_conn ? PQerrorMessage(pgsql_conn) : "NULL connection");
|
||||
if (pgsql_conn) {
|
||||
PQfinish(pgsql_conn);
|
||||
pgsql_conn = NULL;
|
||||
}
|
||||
pthread_mutex_unlock(&conn_lock);
|
||||
return -1;
|
||||
}
|
||||
|
||||
server_version = get_pgsql_version();
|
||||
source_dsn = "pgsql://" + pgsql_user + "@" + pgsql_host + ":" + std::to_string(pgsql_port) + "/" + pgsql_dbname;
|
||||
|
||||
proxy_info("PgSQL_Static_Harvester: Connected to PostgreSQL %s at %s:%d\n",
|
||||
server_version.c_str(), pgsql_host.c_str(), pgsql_port);
|
||||
|
||||
pthread_mutex_unlock(&conn_lock);
|
||||
return 0;
|
||||
}
|
||||
|
||||
void PgSQL_Static_Harvester::disconnect_pgsql() {
|
||||
pthread_mutex_lock(&conn_lock);
|
||||
if (pgsql_conn) {
|
||||
PQfinish(pgsql_conn);
|
||||
pgsql_conn = NULL;
|
||||
}
|
||||
pthread_mutex_unlock(&conn_lock);
|
||||
}
|
||||
|
||||
int PgSQL_Static_Harvester::execute_query(const std::string& query, std::vector<std::vector<std::string>>& results) {
|
||||
if (!pgsql_conn) {
|
||||
proxy_error("PgSQL_Static_Harvester: Not connected to PostgreSQL\n");
|
||||
return -1;
|
||||
}
|
||||
|
||||
PGresult* res = PQexec(pgsql_conn, query.c_str());
|
||||
if (!res) {
|
||||
proxy_error("PgSQL_Static_Harvester: PQexec returned NULL\n");
|
||||
return -1;
|
||||
}
|
||||
ExecStatusType st = PQresultStatus(res);
|
||||
if (st != PGRES_TUPLES_OK && st != PGRES_COMMAND_OK) {
|
||||
proxy_error("PgSQL_Static_Harvester: Query failed: %s\n", PQresultErrorMessage(res));
|
||||
PQclear(res);
|
||||
return -1;
|
||||
}
|
||||
if (st == PGRES_COMMAND_OK) {
|
||||
PQclear(res);
|
||||
return 0;
|
||||
}
|
||||
|
||||
int nrows = PQntuples(res);
|
||||
int ncols = PQnfields(res);
|
||||
results.clear();
|
||||
results.reserve(nrows);
|
||||
for (int r = 0; r < nrows; r++) {
|
||||
std::vector<std::string> row;
|
||||
row.reserve(ncols);
|
||||
for (int c = 0; c < ncols; c++) {
|
||||
row.push_back(PQgetisnull(res, r, c) ? "" : PQgetvalue(res, r, c));
|
||||
}
|
||||
results.push_back(row);
|
||||
}
|
||||
PQclear(res);
|
||||
return 0;
|
||||
}
|
||||
|
||||
std::string PgSQL_Static_Harvester::get_pgsql_version() {
|
||||
std::vector<std::vector<std::string>> rows;
|
||||
if (execute_query("SELECT version();", rows) == 0 && !rows.empty() && !rows[0].empty()) {
|
||||
return rows[0][0];
|
||||
}
|
||||
return "";
|
||||
}
|
||||
|
||||
bool PgSQL_Static_Harvester::is_valid_schema_name(const std::string& name) {
|
||||
if (name.empty()) {
|
||||
return true;
|
||||
}
|
||||
for (char c : name) {
|
||||
if (!std::isalnum(static_cast<unsigned char>(c)) && c != '_' && c != '$') {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
std::string PgSQL_Static_Harvester::escape_sql_string(const std::string& str) {
|
||||
std::string escaped;
|
||||
escaped.reserve(str.length() * 2);
|
||||
for (char c : str) {
|
||||
if (c == '\'') {
|
||||
escaped += "''";
|
||||
} else {
|
||||
escaped += c;
|
||||
}
|
||||
}
|
||||
return escaped;
|
||||
}
|
||||
|
||||
int PgSQL_Static_Harvester::start_run(const std::string& target_id, const std::string& notes) {
|
||||
if (current_run_id >= 0) {
|
||||
proxy_error("PgSQL_Static_Harvester: Run already active (run_id=%d)\n", current_run_id);
|
||||
return -1;
|
||||
}
|
||||
if (connect_pgsql()) {
|
||||
return -1;
|
||||
}
|
||||
|
||||
current_run_id = catalog->create_run(target_id, "pgsql", source_dsn, server_version, notes);
|
||||
if (current_run_id < 0) {
|
||||
proxy_error("PgSQL_Static_Harvester: Failed to create run\n");
|
||||
return -1;
|
||||
}
|
||||
proxy_info("PgSQL_Static_Harvester: Started run_id=%d\n", current_run_id);
|
||||
return current_run_id;
|
||||
}
|
||||
|
||||
int PgSQL_Static_Harvester::finish_run(const std::string& notes) {
|
||||
if (current_run_id < 0) {
|
||||
proxy_error("PgSQL_Static_Harvester: No active run\n");
|
||||
return -1;
|
||||
}
|
||||
int rc = catalog->finish_run(current_run_id, notes);
|
||||
if (rc) {
|
||||
proxy_error("PgSQL_Static_Harvester: Failed to finish run\n");
|
||||
return -1;
|
||||
}
|
||||
proxy_info("PgSQL_Static_Harvester: Finished run_id=%d\n", current_run_id);
|
||||
current_run_id = -1;
|
||||
return 0;
|
||||
}
|
||||
|
||||
int PgSQL_Static_Harvester::harvest_schemas(const std::string& only_schema) {
|
||||
if (current_run_id < 0) {
|
||||
return -1;
|
||||
}
|
||||
if (!is_valid_schema_name(only_schema)) {
|
||||
return -1;
|
||||
}
|
||||
|
||||
std::ostringstream sql;
|
||||
sql << "SELECT schema_name, default_character_set_name, default_collation_name "
|
||||
<< "FROM information_schema.schemata "
|
||||
<< "WHERE schema_name NOT IN ('pg_catalog','information_schema') "
|
||||
<< "AND schema_name NOT LIKE 'pg_toast%'";
|
||||
if (!only_schema.empty()) {
|
||||
sql << " AND schema_name='" << only_schema << "'";
|
||||
}
|
||||
sql << " ORDER BY schema_name";
|
||||
|
||||
std::vector<std::vector<std::string>> rows;
|
||||
if (execute_query(sql.str(), rows)) {
|
||||
return -1;
|
||||
}
|
||||
|
||||
int count = 0;
|
||||
for (const auto& r : rows) {
|
||||
if (r.size() < 3) continue;
|
||||
if (catalog->insert_schema(current_run_id, r[0], r[1], r[2]) >= 0) {
|
||||
count++;
|
||||
}
|
||||
}
|
||||
return count;
|
||||
}
|
||||
|
||||
int PgSQL_Static_Harvester::harvest_objects(const std::string& only_schema) {
|
||||
if (current_run_id < 0) {
|
||||
return -1;
|
||||
}
|
||||
if (!is_valid_schema_name(only_schema)) {
|
||||
return -1;
|
||||
}
|
||||
|
||||
std::vector<std::vector<std::string>> rows;
|
||||
std::ostringstream sql_tables;
|
||||
sql_tables << "SELECT table_schema, table_name, "
|
||||
<< "CASE WHEN table_type='BASE TABLE' THEN 'table' ELSE 'view' END "
|
||||
<< "FROM information_schema.tables "
|
||||
<< "WHERE table_schema NOT IN ('pg_catalog','information_schema') "
|
||||
<< "AND table_schema NOT LIKE 'pg_toast%'";
|
||||
if (!only_schema.empty()) {
|
||||
sql_tables << " AND table_schema='" << only_schema << "'";
|
||||
}
|
||||
if (execute_query(sql_tables.str(), rows)) {
|
||||
return -1;
|
||||
}
|
||||
|
||||
int count = 0;
|
||||
for (const auto& r : rows) {
|
||||
if (r.size() < 3) continue;
|
||||
if (catalog->insert_object(
|
||||
current_run_id, r[0], r[1], r[2],
|
||||
"", 0, 0, 0, "", "", "", ""
|
||||
) >= 0) {
|
||||
count++;
|
||||
}
|
||||
}
|
||||
|
||||
rows.clear();
|
||||
std::ostringstream sql_routines;
|
||||
sql_routines << "SELECT routine_schema, routine_name, 'routine', COALESCE(routine_definition,'') "
|
||||
<< "FROM information_schema.routines "
|
||||
<< "WHERE routine_schema NOT IN ('pg_catalog','information_schema')";
|
||||
if (!only_schema.empty()) {
|
||||
sql_routines << " AND routine_schema='" << only_schema << "'";
|
||||
}
|
||||
if (execute_query(sql_routines.str(), rows) == 0) {
|
||||
for (const auto& r : rows) {
|
||||
if (r.size() < 4) continue;
|
||||
if (catalog->insert_object(
|
||||
current_run_id, r[0], r[1], r[2],
|
||||
"", 0, 0, 0, "", "", "", r[3]
|
||||
) >= 0) {
|
||||
count++;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
rows.clear();
|
||||
std::ostringstream sql_triggers;
|
||||
sql_triggers << "SELECT trigger_schema, trigger_name, 'trigger', COALESCE(action_statement,'') "
|
||||
<< "FROM information_schema.triggers "
|
||||
<< "WHERE trigger_schema NOT IN ('pg_catalog','information_schema')";
|
||||
if (!only_schema.empty()) {
|
||||
sql_triggers << " AND trigger_schema='" << only_schema << "'";
|
||||
}
|
||||
if (execute_query(sql_triggers.str(), rows) == 0) {
|
||||
for (const auto& r : rows) {
|
||||
if (r.size() < 4) continue;
|
||||
if (catalog->insert_object(
|
||||
current_run_id, r[0], r[1], r[2],
|
||||
"", 0, 0, 0, "", "", "", r[3]
|
||||
) >= 0) {
|
||||
count++;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return count;
|
||||
}
|
||||
|
||||
int PgSQL_Static_Harvester::harvest_columns(const std::string& only_schema) {
|
||||
if (current_run_id < 0) return -1;
|
||||
if (!is_valid_schema_name(only_schema)) return -1;
|
||||
|
||||
std::ostringstream sql;
|
||||
sql << "SELECT table_schema, table_name, ordinal_position, column_name, data_type, "
|
||||
<< "COALESCE(udt_name,''), CASE WHEN is_nullable='YES' THEN 1 ELSE 0 END, "
|
||||
<< "COALESCE(column_default,''), '', COALESCE(character_set_name,''), "
|
||||
<< "COALESCE(collation_name,''), COALESCE(column_name,'') "
|
||||
<< "FROM information_schema.columns "
|
||||
<< "WHERE table_schema NOT IN ('pg_catalog','information_schema') "
|
||||
<< "AND table_schema NOT LIKE 'pg_toast%'";
|
||||
if (!only_schema.empty()) {
|
||||
sql << " AND table_schema='" << only_schema << "'";
|
||||
}
|
||||
sql << " ORDER BY table_schema, table_name, ordinal_position";
|
||||
|
||||
std::vector<std::vector<std::string>> rows;
|
||||
if (execute_query(sql.str(), rows)) return -1;
|
||||
|
||||
int count = 0;
|
||||
for (const auto& r : rows) {
|
||||
if (r.size() < 12) continue;
|
||||
char* error = NULL;
|
||||
int cols = 0, affected = 0;
|
||||
SQLite3_result* resultset = NULL;
|
||||
std::ostringstream obj_sql;
|
||||
obj_sql << "SELECT object_id FROM objects "
|
||||
<< "WHERE run_id=" << current_run_id
|
||||
<< " AND schema_name='" << escape_sql_string(r[0]) << "'"
|
||||
<< " AND object_name='" << escape_sql_string(r[1]) << "'"
|
||||
<< " AND object_type IN ('table','view') LIMIT 1;";
|
||||
catalog->get_db()->execute_statement(obj_sql.str().c_str(), &error, &cols, &affected, &resultset);
|
||||
if (error || !resultset || resultset->rows.empty()) {
|
||||
if (error) {
|
||||
free(error);
|
||||
}
|
||||
if (resultset) {
|
||||
delete resultset;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
int object_id = atoi(resultset->rows[0]->fields[0]);
|
||||
delete resultset;
|
||||
|
||||
int is_time = 0;
|
||||
if (r[4] == "date" || r[4] == "time" || r[4] == "timestamp" ||
|
||||
r[4] == "timestamptz" || r[4] == "timetz") {
|
||||
is_time = 1;
|
||||
}
|
||||
int is_id_like = (r[3] == "id" || (r[3].size() > 3 && r[3].substr(r[3].size() - 3) == "_id")) ? 1 : 0;
|
||||
|
||||
if (catalog->insert_column(
|
||||
object_id,
|
||||
atoi(r[2].c_str()), r[3], r[4], r[5], atoi(r[6].c_str()),
|
||||
r[7], r[8], r[9], r[10], r[11],
|
||||
0, 0, 0, is_time, is_id_like
|
||||
) >= 0) {
|
||||
count++;
|
||||
}
|
||||
}
|
||||
catalog->update_object_flags(current_run_id);
|
||||
return count;
|
||||
}
|
||||
|
||||
int PgSQL_Static_Harvester::harvest_indexes(const std::string& only_schema) {
|
||||
if (current_run_id < 0) return -1;
|
||||
if (!is_valid_schema_name(only_schema)) return -1;
|
||||
|
||||
std::ostringstream sql;
|
||||
sql << "SELECT ns.nspname, t.relname, i.relname, "
|
||||
<< "CASE WHEN ix.indisunique THEN 1 ELSE 0 END, "
|
||||
<< "COALESCE(am.amname,''), k.ord, COALESCE(a.attname,''), 0, '', 0 "
|
||||
<< "FROM pg_class t "
|
||||
<< "JOIN pg_namespace ns ON ns.oid=t.relnamespace "
|
||||
<< "JOIN pg_index ix ON ix.indrelid=t.oid "
|
||||
<< "JOIN pg_class i ON i.oid=ix.indexrelid "
|
||||
<< "JOIN pg_am am ON am.oid=i.relam "
|
||||
<< "JOIN LATERAL unnest(ix.indkey) WITH ORDINALITY AS k(attnum, ord) ON TRUE "
|
||||
<< "LEFT JOIN pg_attribute a ON a.attrelid=t.oid AND a.attnum=k.attnum "
|
||||
<< "WHERE t.relkind IN ('r','p') "
|
||||
<< "AND ns.nspname NOT IN ('pg_catalog','information_schema') "
|
||||
<< "AND ns.nspname NOT LIKE 'pg_toast%'";
|
||||
if (!only_schema.empty()) {
|
||||
sql << " AND ns.nspname='" << only_schema << "'";
|
||||
}
|
||||
sql << " ORDER BY ns.nspname, t.relname, i.relname, k.ord";
|
||||
|
||||
std::vector<std::vector<std::string>> rows;
|
||||
if (execute_query(sql.str(), rows)) return -1;
|
||||
|
||||
int count = 0;
|
||||
std::map<std::string, std::vector<std::vector<std::string>>> grouped;
|
||||
for (const auto& r : rows) {
|
||||
if (r.size() < 10) continue;
|
||||
std::string key = r[0] + "." + r[1] + "." + r[2];
|
||||
grouped[key].push_back(r);
|
||||
}
|
||||
|
||||
for (const auto& entry : grouped) {
|
||||
const auto& rows_for_index = entry.second;
|
||||
if (rows_for_index.empty()) continue;
|
||||
const auto& first = rows_for_index[0];
|
||||
|
||||
char* error = NULL;
|
||||
int cols = 0, affected = 0;
|
||||
SQLite3_result* resultset = NULL;
|
||||
std::ostringstream obj_sql;
|
||||
obj_sql << "SELECT object_id FROM objects "
|
||||
<< "WHERE run_id=" << current_run_id
|
||||
<< " AND schema_name='" << escape_sql_string(first[0]) << "'"
|
||||
<< " AND object_name='" << escape_sql_string(first[1]) << "'"
|
||||
<< " AND object_type='table' LIMIT 1;";
|
||||
catalog->get_db()->execute_statement(obj_sql.str().c_str(), &error, &cols, &affected, &resultset);
|
||||
if (error || !resultset || resultset->rows.empty()) {
|
||||
if (error) {
|
||||
free(error);
|
||||
}
|
||||
if (resultset) {
|
||||
delete resultset;
|
||||
}
|
||||
continue;
|
||||
}
|
||||
int object_id = atoi(resultset->rows[0]->fields[0]);
|
||||
delete resultset;
|
||||
|
||||
int is_primary = (first[2] == "PRIMARY" || first[2] == first[1] + "_pkey") ? 1 : 0;
|
||||
int index_id = catalog->insert_index(
|
||||
object_id, first[2], atoi(first[3].c_str()), is_primary, first[4], atol(first[9].c_str())
|
||||
);
|
||||
if (index_id < 0) {
|
||||
continue;
|
||||
}
|
||||
for (const auto& idx_col : rows_for_index) {
|
||||
catalog->insert_index_column(index_id, atoi(idx_col[5].c_str()), idx_col[6], atoi(idx_col[7].c_str()), idx_col[8]);
|
||||
}
|
||||
count++;
|
||||
}
|
||||
|
||||
catalog->update_object_flags(current_run_id);
|
||||
return count;
|
||||
}
|
||||
|
||||
int PgSQL_Static_Harvester::harvest_foreign_keys(const std::string& only_schema) {
|
||||
if (current_run_id < 0) return -1;
|
||||
if (!is_valid_schema_name(only_schema)) return -1;
|
||||
|
||||
std::ostringstream sql;
|
||||
sql << "SELECT tc.table_schema, tc.table_name, tc.constraint_name, "
|
||||
<< "kcu.column_name, ccu.table_schema, ccu.table_name, ccu.column_name, "
|
||||
<< "kcu.ordinal_position, COALESCE(rc.update_rule,''), COALESCE(rc.delete_rule,'') "
|
||||
<< "FROM information_schema.table_constraints tc "
|
||||
<< "JOIN information_schema.key_column_usage kcu "
|
||||
<< "ON tc.constraint_name=kcu.constraint_name "
|
||||
<< "AND tc.table_schema=kcu.table_schema "
|
||||
<< "JOIN information_schema.constraint_column_usage ccu "
|
||||
<< "ON ccu.constraint_name=tc.constraint_name "
|
||||
<< "AND ccu.constraint_schema=tc.constraint_schema "
|
||||
<< "LEFT JOIN information_schema.referential_constraints rc "
|
||||
<< "ON rc.constraint_name=tc.constraint_name "
|
||||
<< "AND rc.constraint_schema=tc.constraint_schema "
|
||||
<< "WHERE tc.constraint_type='FOREIGN KEY' "
|
||||
<< "AND tc.table_schema NOT IN ('pg_catalog','information_schema') "
|
||||
<< "AND tc.table_schema NOT LIKE 'pg_toast%'";
|
||||
if (!only_schema.empty()) {
|
||||
sql << " AND tc.table_schema='" << only_schema << "'";
|
||||
}
|
||||
sql << " ORDER BY tc.table_schema, tc.table_name, tc.constraint_name, kcu.ordinal_position";
|
||||
|
||||
std::vector<std::vector<std::string>> rows;
|
||||
if (execute_query(sql.str(), rows)) return -1;
|
||||
|
||||
int count = 0;
|
||||
int last_fk_id = -1;
|
||||
std::string last_fk_key;
|
||||
for (const auto& r : rows) {
|
||||
if (r.size() < 10) continue;
|
||||
std::string fk_key = r[0] + "." + r[1] + "." + r[2];
|
||||
if (fk_key != last_fk_key) {
|
||||
char* error = NULL;
|
||||
int cols = 0, affected = 0;
|
||||
SQLite3_result* resultset = NULL;
|
||||
std::ostringstream obj_sql;
|
||||
obj_sql << "SELECT object_id FROM objects "
|
||||
<< "WHERE run_id=" << current_run_id
|
||||
<< " AND schema_name='" << escape_sql_string(r[0]) << "'"
|
||||
<< " AND object_name='" << escape_sql_string(r[1]) << "'"
|
||||
<< " AND object_type='table' LIMIT 1;";
|
||||
catalog->get_db()->execute_statement(obj_sql.str().c_str(), &error, &cols, &affected, &resultset);
|
||||
if (error || !resultset || resultset->rows.empty()) {
|
||||
if (error) {
|
||||
free(error);
|
||||
}
|
||||
if (resultset) {
|
||||
delete resultset;
|
||||
}
|
||||
last_fk_id = -1;
|
||||
last_fk_key.clear();
|
||||
continue;
|
||||
}
|
||||
int child_object_id = atoi(resultset->rows[0]->fields[0]);
|
||||
delete resultset;
|
||||
last_fk_id = catalog->insert_foreign_key(
|
||||
current_run_id, child_object_id, r[2], r[4], r[5], r[8], r[9]
|
||||
);
|
||||
last_fk_key = fk_key;
|
||||
}
|
||||
if (last_fk_id >= 0) {
|
||||
catalog->insert_foreign_key_column(last_fk_id, atoi(r[7].c_str()), r[3], r[6]);
|
||||
count++;
|
||||
}
|
||||
}
|
||||
|
||||
catalog->update_object_flags(current_run_id);
|
||||
return count;
|
||||
}
|
||||
|
||||
int PgSQL_Static_Harvester::harvest_view_definitions(const std::string& only_schema) {
|
||||
if (current_run_id < 0) return -1;
|
||||
if (!is_valid_schema_name(only_schema)) return -1;
|
||||
|
||||
std::ostringstream sql;
|
||||
sql << "SELECT schemaname, viewname, definition FROM pg_views "
|
||||
<< "WHERE schemaname NOT IN ('pg_catalog','information_schema') "
|
||||
<< "AND schemaname NOT LIKE 'pg_toast%'";
|
||||
if (!only_schema.empty()) {
|
||||
sql << " AND schemaname='" << only_schema << "'";
|
||||
}
|
||||
|
||||
std::vector<std::vector<std::string>> rows;
|
||||
if (execute_query(sql.str(), rows)) return -1;
|
||||
|
||||
int count = 0;
|
||||
for (const auto& row : rows) {
|
||||
if (row.size() < 3) continue;
|
||||
char* error = NULL;
|
||||
int cols = 0, affected = 0;
|
||||
std::ostringstream update_sql;
|
||||
update_sql << "UPDATE objects SET definition_sql = '" << escape_sql_string(row[2]) << "' "
|
||||
<< "WHERE run_id = " << current_run_id
|
||||
<< " AND schema_name = '" << escape_sql_string(row[0]) << "'"
|
||||
<< " AND object_name = '" << escape_sql_string(row[1]) << "'"
|
||||
<< " AND object_type = 'view';";
|
||||
catalog->get_db()->execute_statement(update_sql.str().c_str(), &error, &cols, &affected);
|
||||
if (error) {
|
||||
free(error);
|
||||
}
|
||||
if (affected > 0) {
|
||||
count++;
|
||||
}
|
||||
}
|
||||
return count;
|
||||
}
|
||||
|
||||
int PgSQL_Static_Harvester::build_quick_profiles() {
|
||||
if (current_run_id < 0) return -1;
|
||||
|
||||
char* error = NULL;
|
||||
int cols = 0, affected = 0;
|
||||
SQLite3_result* resultset = NULL;
|
||||
std::ostringstream sql;
|
||||
sql << "SELECT object_id, object_name, table_rows_est, data_length, index_length, "
|
||||
"has_primary_key, has_foreign_keys, has_time_column "
|
||||
"FROM objects WHERE run_id=" << current_run_id << " AND object_type='table'";
|
||||
catalog->get_db()->execute_statement(sql.str().c_str(), &error, &cols, &affected, &resultset);
|
||||
if (error) {
|
||||
free(error);
|
||||
if (resultset) delete resultset;
|
||||
return -1;
|
||||
}
|
||||
|
||||
int count = 0;
|
||||
if (resultset) {
|
||||
for (auto row : resultset->rows) {
|
||||
if (row->cnt < 8) continue;
|
||||
const std::string name = row->fields[1] ? row->fields[1] : "";
|
||||
std::string kind = "unknown";
|
||||
std::string lname = name;
|
||||
std::transform(lname.begin(), lname.end(), lname.begin(), ::tolower);
|
||||
if (lname.find("log") != std::string::npos || lname.find("event") != std::string::npos || lname.find("audit") != std::string::npos) {
|
||||
kind = "log";
|
||||
} else if (lname.find("order") != std::string::npos || lname.find("invoice") != std::string::npos || lname.find("payment") != std::string::npos || lname.find("transaction") != std::string::npos) {
|
||||
kind = "fact";
|
||||
} else if (lname.find("user") != std::string::npos || lname.find("customer") != std::string::npos || lname.find("account") != std::string::npos || lname.find("product") != std::string::npos) {
|
||||
kind = "entity";
|
||||
}
|
||||
|
||||
json p;
|
||||
p["guessed_kind"] = kind;
|
||||
p["rows_est"] = row->fields[2] ? atoll(row->fields[2]) : 0;
|
||||
p["size_bytes"] = (row->fields[3] ? atoll(row->fields[3]) : 0) + (row->fields[4] ? atoll(row->fields[4]) : 0);
|
||||
p["engine"] = "pgsql";
|
||||
p["has_primary_key"] = row->fields[5] ? atoi(row->fields[5]) : 0;
|
||||
p["has_foreign_keys"] = row->fields[6] ? atoi(row->fields[6]) : 0;
|
||||
p["has_time_column"] = row->fields[7] ? atoi(row->fields[7]) : 0;
|
||||
|
||||
if (catalog->upsert_profile(current_run_id, atoi(row->fields[0]), "table_quick", p.dump()) == 0) {
|
||||
count++;
|
||||
}
|
||||
}
|
||||
delete resultset;
|
||||
}
|
||||
return count;
|
||||
}
|
||||
|
||||
int PgSQL_Static_Harvester::rebuild_fts_index() {
|
||||
if (current_run_id < 0) return -1;
|
||||
return catalog->rebuild_fts_index(current_run_id);
|
||||
}
|
||||
|
||||
int PgSQL_Static_Harvester::run_full_harvest(const std::string& target_id, const std::string& only_schema, const std::string& notes) {
|
||||
if (start_run(target_id, notes) < 0) return -1;
|
||||
if (harvest_schemas(only_schema) < 0) { finish_run("Failed during schema harvest"); return -1; }
|
||||
if (harvest_objects(only_schema) < 0) { finish_run("Failed during object harvest"); return -1; }
|
||||
if (harvest_columns(only_schema) < 0) { finish_run("Failed during column harvest"); return -1; }
|
||||
if (harvest_indexes(only_schema) < 0) { finish_run("Failed during index harvest"); return -1; }
|
||||
if (harvest_foreign_keys(only_schema) < 0) { finish_run("Failed during foreign key harvest"); return -1; }
|
||||
if (harvest_view_definitions(only_schema) < 0) { finish_run("Failed during view definition harvest"); return -1; }
|
||||
if (build_quick_profiles() < 0) { finish_run("Failed during profile building"); return -1; }
|
||||
if (rebuild_fts_index() < 0) { finish_run("Failed during FTS rebuild"); return -1; }
|
||||
int final_run_id = current_run_id;
|
||||
finish_run("Harvest completed successfully");
|
||||
return final_run_id;
|
||||
}
|
||||
|
||||
std::string PgSQL_Static_Harvester::get_harvest_stats() {
|
||||
if (current_run_id < 0) {
|
||||
return "{\"error\": \"No active run\"}";
|
||||
}
|
||||
return get_harvest_stats(current_run_id);
|
||||
}
|
||||
|
||||
std::string PgSQL_Static_Harvester::get_harvest_stats(int run_id) {
|
||||
char* error = NULL;
|
||||
int cols = 0, affected = 0;
|
||||
SQLite3_result* resultset = NULL;
|
||||
json stats;
|
||||
stats["run_id"] = run_id;
|
||||
stats["protocol"] = "pgsql";
|
||||
|
||||
std::ostringstream q1;
|
||||
q1 << "SELECT object_type, COUNT(*) FROM objects WHERE run_id=" << run_id << " GROUP BY object_type";
|
||||
catalog->get_db()->execute_statement(q1.str().c_str(), &error, &cols, &affected, &resultset);
|
||||
json objects = json::object();
|
||||
if (!error && resultset) {
|
||||
for (auto row : resultset->rows) {
|
||||
if (row->cnt >= 2) {
|
||||
objects[row->fields[0] ? row->fields[0] : "unknown"] = row->fields[1] ? atoi(row->fields[1]) : 0;
|
||||
}
|
||||
}
|
||||
delete resultset;
|
||||
} else if (error) {
|
||||
free(error);
|
||||
error = NULL;
|
||||
if (resultset) delete resultset;
|
||||
}
|
||||
stats["objects"] = objects;
|
||||
|
||||
auto get_count = [&](const std::string& table) -> int {
|
||||
char* e = NULL;
|
||||
int c = 0, a = 0;
|
||||
SQLite3_result* rs = NULL;
|
||||
std::ostringstream q;
|
||||
q << "SELECT COUNT(*) FROM " << table << " WHERE run_id=" << run_id;
|
||||
catalog->get_db()->execute_statement(q.str().c_str(), &e, &c, &a, &rs);
|
||||
int out = 0;
|
||||
if (!e && rs && !rs->rows.empty() && rs->rows[0]->cnt > 0 && rs->rows[0]->fields[0]) {
|
||||
out = atoi(rs->rows[0]->fields[0]);
|
||||
}
|
||||
if (e) free(e);
|
||||
if (rs) delete rs;
|
||||
return out;
|
||||
};
|
||||
|
||||
stats["columns"] = get_count("columns");
|
||||
stats["foreign_keys"] = get_count("foreign_keys");
|
||||
|
||||
return stats.dump();
|
||||
}
|
||||
|
||||
#endif /* PROXYSQLGENAI */
|
||||
Loading…
Reference in new issue