@neverinfamous/postgres-mcp 1.1.0 → 1.3.0

package/dist/adapters/postgresql/tools/stats/advanced.js

@@ -5,6 +5,7 @@
  */
 import { readOnly } from "../../../../utils/annotations.js";
 import { getToolIcons } from "../../../../utils/icons.js";
+import { formatPostgresError } from "../core/error-helpers.js";
 import { 
 // Base schemas for MCP visibility
 StatsTimeSeriesSchemaBase, StatsDistributionSchemaBase, StatsHypothesisSchemaBase, StatsSamplingSchemaBase, 
@@ -236,115 +237,118 @@ export function createStatsTimeSeriesTool(adapter) {
         annotations: readOnly("Time Series Analysis"),
         icons: getToolIcons("stats", readOnly("Time Series Analysis")),
         handler: async (params, _context) => {
-            const { table, valueColumn, timeColumn, interval, aggregation, schema, where, params: queryParams, limit, groupBy, groupLimit, } = StatsTimeSeriesSchema.parse(params);
-            const schemaPrefix = schema ? `"${schema}".` : "";
-            const whereClause = where ? `WHERE ${where}` : "";
-            const agg = aggregation ?? "avg";
-            // Handle limit: undefined uses default (100), 0 means no limit
-            // Track whether user explicitly provided a limit
-            const userProvidedLimit = limit !== undefined;
-            const DEFAULT_LIMIT = 100;
-            // limit === 0 means "no limit", otherwise use provided limit or default
-            const effectiveLimit = limit === 0 ? undefined : (limit ?? DEFAULT_LIMIT);
-            const usingDefaultLimit = !userProvidedLimit && effectiveLimit !== undefined;
-            // First check if table exists
-            const schemaName = schema ?? "public";
-            const tableCheckQuery = `
+            try {
+                const { table, valueColumn, timeColumn, interval, aggregation, schema, where, params: queryParams, limit, groupBy, groupLimit, } = StatsTimeSeriesSchema.parse(params);
+                const schemaPrefix = schema ? `"${schema}".` : "";
+                const whereClause = where ? `WHERE ${where}` : "";
+                const agg = aggregation ?? "avg";
+                // Handle limit: undefined uses default (100), 0 means no limit
+                // Track whether user explicitly provided a limit
+                const userProvidedLimit = limit !== undefined;
+                const DEFAULT_LIMIT = 100;
+                // limit === 0 means "no limit", otherwise use provided limit or default
+                const effectiveLimit = limit === 0 ? undefined : (limit ?? DEFAULT_LIMIT);
+                const usingDefaultLimit = !userProvidedLimit && effectiveLimit !== undefined;
+                // First check if table exists
+                const schemaName = schema ?? "public";
+                const tableCheckQuery = `
         SELECT 1 FROM information_schema.tables 
         WHERE table_schema = '${schemaName}' AND table_name = '${table}'
       `;
-            const tableCheckResult = await adapter.executeQuery(tableCheckQuery);
-            if (tableCheckResult.rows?.length === 0) {
-                throw new Error(`Table "${schemaName}.${table}" not found`);
-            }
-            // Validate timeColumn is a timestamp/date type
-            const typeCheckQuery = `
+                const tableCheckResult = await adapter.executeQuery(tableCheckQuery);
+                if (tableCheckResult.rows?.length === 0) {
+                    throw new Error(`Table "${schemaName}.${table}" not found`);
+                }
+                // Validate timeColumn is a timestamp/date type
+                const typeCheckQuery = `
                 SELECT data_type 
                 FROM information_schema.columns 
                 WHERE table_schema = '${schemaName}' 
                 AND table_name = '${table}'
                 AND column_name = '${timeColumn}'
             `;
-            const typeResult = await adapter.executeQuery(typeCheckQuery);
-            const typeRow = typeResult.rows?.[0];
-            if (!typeRow) {
-                throw new Error(`Column "${timeColumn}" not found in table "${schemaName}.${table}"`);
-            }
-            const validTypes = [
-                "timestamp without time zone",
-                "timestamp with time zone",
-                "date",
-                "time",
-                "time without time zone",
-                "time with time zone",
-            ];
-            if (!validTypes.includes(typeRow.data_type)) {
-                throw new Error(`Column "${timeColumn}" is type "${typeRow.data_type}" but must be a timestamp or date type for time series analysis`);
-            }
-            // Note: schemaName already defined above for table check
-            // Validate valueColumn exists and is numeric
-            const numericTypes = [
-                "integer",
-                "bigint",
-                "smallint",
-                "numeric",
-                "decimal",
-                "real",
-                "double precision",
-                "money",
-            ];
-            const valueTypeQuery = `
+                const typeResult = await adapter.executeQuery(typeCheckQuery);
+                const typeRow = typeResult.rows?.[0];
+                if (!typeRow) {
+                    throw new Error(`Column "${timeColumn}" not found in table "${schemaName}.${table}"`);
+                }
+                const validTypes = [
+                    "timestamp without time zone",
+                    "timestamp with time zone",
+                    "date",
+                    "time",
+                    "time without time zone",
+                    "time with time zone",
+                ];
+                if (!validTypes.includes(typeRow.data_type)) {
+                    throw new Error(`Column "${timeColumn}" is type "${typeRow.data_type}" but must be a timestamp or date type for time series analysis`);
+                }
+                // Note: schemaName already defined above for table check
+                // Validate valueColumn exists and is numeric
+                const numericTypes = [
+                    "integer",
+                    "bigint",
+                    "smallint",
+                    "numeric",
+                    "decimal",
+                    "real",
+                    "double precision",
+                    "money",
+                ];
+                const valueTypeQuery = `
         SELECT data_type 
         FROM information_schema.columns 
         WHERE table_schema = '${schemaName}' 
         AND table_name = '${table}'
         AND column_name = '${valueColumn}'
       `;
-            const valueTypeResult = await adapter.executeQuery(valueTypeQuery);
-            const valueTypeRow = valueTypeResult.rows?.[0];
-            if (!valueTypeRow) {
-                throw new Error(`Column "${valueColumn}" not found in table "${schemaName}.${table}"`);
-            }
-            if (!numericTypes.includes(valueTypeRow.data_type)) {
-                throw new Error(`Column "${valueColumn}" is type "${valueTypeRow.data_type}" but must be a numeric type for time series aggregation`);
-            }
-            // Helper to map bucket row - convert Date to ISO string for JSON Schema
-            // Handles both Date objects (from real DB) and strings (from mocks)
-            const mapBucket = (row) => {
-                const timeBucketValue = row["time_bucket"];
-                let timeBucket;
-                if (timeBucketValue instanceof Date) {
-                    timeBucket = timeBucketValue.toISOString();
+                const valueTypeResult = await adapter.executeQuery(valueTypeQuery);
+                const valueTypeRow = valueTypeResult.rows?.[0];
+                if (!valueTypeRow) {
+                    throw new Error(`Column "${valueColumn}" not found in table "${schemaName}.${table}"`);
                 }
-                else if (typeof timeBucketValue === "string") {
-                    timeBucket = timeBucketValue;
-                }
-                else {
-                    // Fallback: null, undefined, or unexpected type
-                    timeBucket = "";
+                if (!numericTypes.includes(valueTypeRow.data_type)) {
+                    throw new Error(`Column "${valueColumn}" is type "${valueTypeRow.data_type}" but must be a numeric type for time series aggregation`);
                 }
-                return {
-                    timeBucket,
-                    value: Number(row["value"]),
-                    count: Number(row["count"]),
+                // Helper to map bucket row - convert Date to ISO string for JSON Schema
+                // Handles both Date objects (from real DB) and strings (from mocks)
+                const mapBucket = (row) => {
+                    const timeBucketValue = row["time_bucket"];
+                    let timeBucket;
+                    if (timeBucketValue instanceof Date) {
+                        timeBucket = timeBucketValue.toISOString();
+                    }
+                    else if (typeof timeBucketValue === "string") {
+                        timeBucket = timeBucketValue;
+                    }
+                    else {
+                        // Fallback: null, undefined, or unexpected type
+                        timeBucket = "";
+                    }
+                    return {
+                        timeBucket,
+                        value: Number(row["value"]),
+                        count: Number(row["count"]),
+                    };
                 };
-            };
-            if (groupBy !== undefined) {
-                // Handle groupLimit: undefined uses default (20), 0 means no limit
-                const DEFAULT_GROUP_LIMIT = 20;
-                const userProvidedGroupLimit = groupLimit !== undefined;
-                const effectiveGroupLimit = groupLimit === 0 ? undefined : (groupLimit ?? DEFAULT_GROUP_LIMIT);
-                // First get total count of distinct groups for truncation indicator
-                const groupCountSql = `
-          SELECT COUNT(DISTINCT "${groupBy}") as total_groups
+                if (groupBy !== undefined) {
+                    // Handle groupLimit: undefined uses default (20), 0 means no limit
+                    const DEFAULT_GROUP_LIMIT = 20;
+                    const userProvidedGroupLimit = groupLimit !== undefined;
+                    const effectiveGroupLimit = groupLimit === 0 ? undefined : (groupLimit ?? DEFAULT_GROUP_LIMIT);
+                    // First get total count of distinct groups for truncation indicator
+                    // COUNT(DISTINCT) excludes NULLs per SQL standard, so add 1 if any NULLs exist
+                    const groupCountSql = `
+          SELECT COUNT(DISTINCT "${groupBy}") + 
+            CASE WHEN COUNT(*) > COUNT("${groupBy}") THEN 1 ELSE 0 END as total_groups
           FROM ${schemaPrefix}"${table}"
           ${whereClause}
         `;
-                const groupCountResult = await adapter.executeQuery(groupCountSql);
-                const totalGroupCount = Number(groupCountResult.rows?.[0]
-                    ?.total_groups ?? 0);
-                // Grouped time series
-                const sql = `
+                    const groupCountResult = await adapter.executeQuery(groupCountSql);
+                    const totalGroupCount = Number(groupCountResult.rows?.[0]
+                        ?.total_groups ?? 0);
+                    // Grouped time series
+                    const sql = `
                     SELECT 
                         "${groupBy}" as group_key,
                         DATE_TRUNC('${interval}', "${timeColumn}") as time_bucket,
@@ -355,77 +359,78 @@ export function createStatsTimeSeriesTool(adapter) {
                     GROUP BY "${groupBy}", DATE_TRUNC('${interval}', "${timeColumn}")
                     ORDER BY "${groupBy}", time_bucket DESC
                 `;
-                const result = await adapter.executeQuery(sql, ...(queryParams !== undefined && queryParams.length > 0
-                    ? [queryParams]
-                    : []));
-                const rows = result.rows ?? [];
-                // Group results by group_key
-                const groupsMap = new Map();
-                const groupsTotalCount = new Map();
-                let groupsProcessed = 0;
-                for (const row of rows) {
-                    const key = row["group_key"];
-                    if (!groupsMap.has(key)) {
-                        // Check if we've hit the group limit
-                        if (effectiveGroupLimit !== undefined &&
-                            groupsProcessed >= effectiveGroupLimit) {
-                            continue;
+                    const result = await adapter.executeQuery(sql, ...(queryParams !== undefined && queryParams.length > 0
+                        ? [queryParams]
+                        : []));
+                    const rows = result.rows ?? [];
+                    // Group results by group_key
+                    const groupsMap = new Map();
+                    const groupsTotalCount = new Map();
+                    let groupsProcessed = 0;
+                    for (const row of rows) {
+                        const key = row["group_key"];
+                        if (!groupsMap.has(key)) {
+                            // Check if we've hit the group limit
+                            if (effectiveGroupLimit !== undefined &&
+                                groupsProcessed >= effectiveGroupLimit) {
+                                continue;
+                            }
+                            groupsMap.set(key, []);
+                            groupsTotalCount.set(key, 0);
+                            groupsProcessed++;
+                        }
+                        const currentTotal = groupsTotalCount.get(key) ?? 0;
+                        groupsTotalCount.set(key, currentTotal + 1);
+                        const bucketList = groupsMap.get(key);
+                        // Only add if no limit or under limit
+                        if (bucketList !== undefined &&
+                            (effectiveLimit === undefined ||
+                                bucketList.length < effectiveLimit)) {
+                            bucketList.push(mapBucket(row));
                         }
-                        groupsMap.set(key, []);
-                        groupsTotalCount.set(key, 0);
-                        groupsProcessed++;
                     }
-                    const currentTotal = groupsTotalCount.get(key) ?? 0;
-                    groupsTotalCount.set(key, currentTotal + 1);
-                    const bucketList = groupsMap.get(key);
-                    // Only add if no limit or under limit
-                    if (bucketList !== undefined &&
-                        (effectiveLimit === undefined || bucketList.length < effectiveLimit)) {
-                        bucketList.push(mapBucket(row));
+                    const groups = Array.from(groupsMap.entries()).map(([key, buckets]) => ({
+                        groupKey: key,
+                        buckets,
+                    }));
+                    // Build response with truncation indicators
+                    const response = {
+                        table: `${schema ?? "public"}.${table}`,
+                        valueColumn,
+                        timeColumn,
+                        interval,
+                        aggregation: agg,
+                        groupBy,
+                        groups,
+                        count: groups.length,
+                    };
+                    // Add truncation indicators when groups are limited
+                    const groupsTruncated = effectiveGroupLimit !== undefined &&
+                        totalGroupCount > effectiveGroupLimit;
+                    if (groupsTruncated || !userProvidedGroupLimit) {
+                        response["truncated"] = groupsTruncated;
+                        response["totalGroupCount"] = totalGroupCount;
                     }
+                    return response;
                 }
-                const groups = Array.from(groupsMap.entries()).map(([key, buckets]) => ({
-                    groupKey: key,
-                    buckets,
-                }));
-                // Build response with truncation indicators
-                const response = {
-                    table: `${schema ?? "public"}.${table}`,
-                    valueColumn,
-                    timeColumn,
-                    interval,
-                    aggregation: agg,
-                    groupBy,
-                    groups,
-                    count: groups.length,
-                };
-                // Add truncation indicators when groups are limited
-                const groupsTruncated = effectiveGroupLimit !== undefined &&
-                    totalGroupCount > effectiveGroupLimit;
-                if (groupsTruncated || !userProvidedGroupLimit) {
-                    response["truncated"] = groupsTruncated;
-                    response["totalGroupCount"] = totalGroupCount;
-                }
-                return response;
-            }
-            // Ungrouped time series
-            // Build LIMIT clause: no LIMIT if effectiveLimit is undefined (limit: 0)
-            const limitClause = effectiveLimit !== undefined ? `LIMIT ${String(effectiveLimit)}` : "";
-            // Get total count if using default limit (for truncation indicator)
-            let totalCount;
-            if (usingDefaultLimit) {
-                const countSql = `
+                // Ungrouped time series
+                // Build LIMIT clause: no LIMIT if effectiveLimit is undefined (limit: 0)
+                const limitClause = effectiveLimit !== undefined ? `LIMIT ${String(effectiveLimit)}` : "";
+                // Get total count if using default limit (for truncation indicator)
+                let totalCount;
+                if (usingDefaultLimit) {
+                    const countSql = `
           SELECT COUNT(DISTINCT DATE_TRUNC('${interval}', "${timeColumn}")) as total_buckets
           FROM ${schemaPrefix}"${table}"
           ${whereClause}
         `;
-                const countResult = await adapter.executeQuery(countSql, ...(queryParams !== undefined && queryParams.length > 0
-                    ? [queryParams]
-                    : []));
-                const countRow = countResult.rows?.[0];
-                totalCount = countRow ? Number(countRow.total_buckets) : undefined;
-            }
-            const sql = `
+                    const countResult = await adapter.executeQuery(countSql, ...(queryParams !== undefined && queryParams.length > 0
+                        ? [queryParams]
+                        : []));
+                    const countRow = countResult.rows?.[0];
+                    totalCount = countRow ? Number(countRow.total_buckets) : undefined;
+                }
+                const sql = `
                 SELECT 
                     DATE_TRUNC('${interval}', "${timeColumn}") as time_bucket,
                     ${agg.toUpperCase()}("${valueColumn}")::numeric(20,6) as value,
@@ -436,25 +441,32 @@ export function createStatsTimeSeriesTool(adapter) {
                 ORDER BY time_bucket DESC
                 ${limitClause}
             `;
-            const result = await adapter.executeQuery(sql, ...(queryParams !== undefined && queryParams.length > 0
-                ? [queryParams]
-                : []));
-            const buckets = (result.rows ?? []).map((row) => mapBucket(row));
-            // Build response
-            const response = {
-                table: `${schema ?? "public"}.${table}`,
-                valueColumn,
-                timeColumn,
-                interval,
-                aggregation: agg,
-                buckets,
-            };
-            // Add truncation indicators when default limit was applied
-            if (usingDefaultLimit && totalCount !== undefined) {
-                response["truncated"] = buckets.length < totalCount;
-                response["totalCount"] = totalCount;
+                const result = await adapter.executeQuery(sql, ...(queryParams !== undefined && queryParams.length > 0
+                    ? [queryParams]
+                    : []));
+                const buckets = (result.rows ?? []).map((row) => mapBucket(row));
+                // Build response
+                const response = {
+                    table: `${schema ?? "public"}.${table}`,
+                    valueColumn,
+                    timeColumn,
+                    interval,
+                    aggregation: agg,
+                    buckets,
+                };
+                // Add truncation indicators when default limit was applied
+                if (usingDefaultLimit && totalCount !== undefined) {
+                    response["truncated"] = buckets.length < totalCount;
+                    response["totalCount"] = totalCount;
+                }
+                return response;
+            }
+            catch (error) {
+                return {
+                    success: false,
+                    error: formatPostgresError(error, { tool: "pg_stats_time_series" }),
+                };
             }
-            return response;
         },
     };
 }
@@ -471,22 +483,23 @@ export function createStatsDistributionTool(adapter) {
         annotations: readOnly("Distribution Analysis"),
         icons: getToolIcons("stats", readOnly("Distribution Analysis")),
         handler: async (params, _context) => {
-            const parsed = StatsDistributionSchema.parse(params);
-            const { table, column, buckets, schema, where, params: queryParams, groupBy, groupLimit, } = parsed;
-            const schemaName = schema ?? "public";
-            const schemaPrefix = schema ? `"${schema}".` : "";
-            const whereClause = where ? `WHERE ${where}` : "";
-            const numBuckets = buckets ?? 10;
-            // Validate column exists and is numeric
-            await validateNumericColumn(adapter, table, column, schemaName);
-            // Helper to compute skewness and kurtosis for a given group
-            const computeMoments = async (groupFilter) => {
-                const filterClause = groupFilter
-                    ? whereClause
-                        ? `${whereClause} AND ${groupFilter}`
-                        : `WHERE ${groupFilter}`
-                    : whereClause;
-                const statsQuery = `
+            try {
+                const parsed = StatsDistributionSchema.parse(params);
+                const { table, column, buckets, schema, where, params: queryParams, groupBy, groupLimit, } = parsed;
+                const schemaName = schema ?? "public";
+                const schemaPrefix = schema ? `"${schema}".` : "";
+                const whereClause = where ? `WHERE ${where}` : "";
+                const numBuckets = buckets ?? 10;
+                // Validate column exists and is numeric
+                await validateNumericColumn(adapter, table, column, schemaName);
+                // Helper to compute skewness and kurtosis for a given group
+                const computeMoments = async (groupFilter) => {
+                    const filterClause = groupFilter
+                        ? whereClause
+                            ? `${whereClause} AND ${groupFilter}`
+                            : `WHERE ${groupFilter}`
+                        : whereClause;
+                    const statsQuery = `
                     WITH stats AS (
                         SELECT 
                             MIN("${column}") as min_val,
@@ -516,28 +529,28 @@ export function createStatsDistributionTool(adapter) {
                     )
                     SELECT * FROM moments
                 `;
-                const result = await adapter.executeQuery(statsQuery, ...(queryParams !== undefined && queryParams.length > 0
-                    ? [queryParams]
-                    : []));
-                const row = result.rows?.[0];
-                if (row?.["min_val"] == null || row["max_val"] == null) {
-                    return null;
-                }
-                return {
-                    minVal: Number(row["min_val"]),
-                    maxVal: Number(row["max_val"]),
-                    skewness: row["skewness"] !== null ? Number(row["skewness"]) : null,
-                    kurtosis: row["kurtosis"] !== null ? Number(row["kurtosis"]) : null,
+                    const result = await adapter.executeQuery(statsQuery, ...(queryParams !== undefined && queryParams.length > 0
+                        ? [queryParams]
+                        : []));
+                    const row = result.rows?.[0];
+                    if (row?.["min_val"] == null || row["max_val"] == null) {
+                        return null;
+                    }
+                    return {
+                        minVal: Number(row["min_val"]),
+                        maxVal: Number(row["max_val"]),
+                        skewness: row["skewness"] !== null ? Number(row["skewness"]) : null,
+                        kurtosis: row["kurtosis"] !== null ? Number(row["kurtosis"]) : null,
+                    };
                 };
-            };
-            // Helper to generate histogram for given min/max
-            const generateHistogram = async (minVal, maxVal, groupFilter) => {
-                const filterClause = groupFilter
-                    ? whereClause
-                        ? `${whereClause} AND ${groupFilter}`
-                        : `WHERE ${groupFilter}`
-                    : whereClause;
-                const histogramQuery = `
+                // Helper to generate histogram for given min/max
+                const generateHistogram = async (minVal, maxVal, groupFilter) => {
+                    const filterClause = groupFilter
+                        ? whereClause
+                            ? `${whereClause} AND ${groupFilter}`
+                            : `WHERE ${groupFilter}`
+                        : whereClause;
+                    const histogramQuery = `
                     SELECT 
                         WIDTH_BUCKET("${column}", ${String(minVal)}, ${String(maxVal + 0.0001)}, ${String(numBuckets)}) as bucket,
                         COUNT(*) as frequency,
@@ -548,92 +561,99 @@ export function createStatsDistributionTool(adapter) {
                     GROUP BY WIDTH_BUCKET("${column}", ${String(minVal)}, ${String(maxVal + 0.0001)}, ${String(numBuckets)})
                     ORDER BY bucket
                 `;
-                const result = await adapter.executeQuery(histogramQuery, ...(queryParams !== undefined && queryParams.length > 0
-                    ? [queryParams]
-                    : []));
-                return (result.rows ?? []).map((row) => ({
-                    bucket: Number(row["bucket"]),
-                    frequency: Number(row["frequency"]),
-                    rangeMin: Number(row["bucket_min"]),
-                    rangeMax: Number(row["bucket_max"]),
-                }));
-            };
-            if (groupBy !== undefined) {
-                // Handle groupLimit: undefined uses default (20), 0 means no limit
-                const DEFAULT_GROUP_LIMIT = 20;
-                const userProvidedGroupLimit = groupLimit !== undefined;
-                const effectiveGroupLimit = groupLimit === 0 ? undefined : (groupLimit ?? DEFAULT_GROUP_LIMIT);
-                // Get distinct groups first
-                const groupsQuery = `
+                    const result = await adapter.executeQuery(histogramQuery, ...(queryParams !== undefined && queryParams.length > 0
+                        ? [queryParams]
+                        : []));
+                    return (result.rows ?? []).map((row) => ({
+                        bucket: Number(row["bucket"]),
+                        frequency: Number(row["frequency"]),
+                        rangeMin: Number(row["bucket_min"]),
+                        rangeMax: Number(row["bucket_max"]),
+                    }));
+                };
+                if (groupBy !== undefined) {
+                    // Handle groupLimit: undefined uses default (20), 0 means no limit
+                    const DEFAULT_GROUP_LIMIT = 20;
+                    const userProvidedGroupLimit = groupLimit !== undefined;
+                    const effectiveGroupLimit = groupLimit === 0 ? undefined : (groupLimit ?? DEFAULT_GROUP_LIMIT);
+                    // Get distinct groups first
+                    const groupsQuery = `
                     SELECT DISTINCT "${groupBy}" as group_key
                     FROM ${schemaPrefix}"${table}"
                     ${whereClause}
                     ORDER BY "${groupBy}"
                 `;
-                const groupsResult = await adapter.executeQuery(groupsQuery, ...(queryParams !== undefined && queryParams.length > 0
-                    ? [queryParams]
-                    : []));
-                const allGroupKeys = (groupsResult.rows ?? []).map((r) => r["group_key"]);
-                const totalGroupCount = allGroupKeys.length;
-                // Apply group limit
-                const groupKeys = effectiveGroupLimit !== undefined
-                    ? allGroupKeys.slice(0, effectiveGroupLimit)
-                    : allGroupKeys;
-                // Process each group
-                const groups = [];
-                for (const groupKey of groupKeys) {
-                    const groupFilter = typeof groupKey === "string"
-                        ? `"${groupBy}" = '${groupKey.replace(/'/g, "''")}'`
-                        : `"${groupBy}" = ${String(groupKey)}`;
-                    const moments = await computeMoments(groupFilter);
-                    if (moments === null)
-                        continue;
-                    const { minVal, maxVal, skewness, kurtosis } = moments;
-                    const bucketWidth = Math.round(((maxVal - minVal) / numBuckets) * 1e6) / 1e6;
-                    const histogram = await generateHistogram(minVal, maxVal, groupFilter);
-                    groups.push({
-                        groupKey,
-                        range: { min: minVal, max: maxVal },
-                        bucketWidth,
-                        skewness,
-                        kurtosis,
-                        histogram,
-                    });
+                    const groupsResult = await adapter.executeQuery(groupsQuery, ...(queryParams !== undefined && queryParams.length > 0
+                        ? [queryParams]
+                        : []));
+                    const allGroupKeys = (groupsResult.rows ?? []).map((r) => r["group_key"]);
+                    const totalGroupCount = allGroupKeys.length;
+                    // Apply group limit
+                    const groupKeys = effectiveGroupLimit !== undefined
+                        ? allGroupKeys.slice(0, effectiveGroupLimit)
+                        : allGroupKeys;
+                    // Process each group
+                    const groups = [];
+                    for (const groupKey of groupKeys) {
+                        const groupFilter = typeof groupKey === "string"
+                            ? `"${groupBy}" = '${groupKey.replace(/'/g, "''")}'`
+                            : `"${groupBy}" = ${String(groupKey)}`;
+                        const moments = await computeMoments(groupFilter);
+                        if (moments === null)
+                            continue;
+                        const { minVal, maxVal, skewness, kurtosis } = moments;
+                        const bucketWidth = Math.round(((maxVal - minVal) / numBuckets) * 1e6) / 1e6;
+                        const histogram = await generateHistogram(minVal, maxVal, groupFilter);
+                        groups.push({
+                            groupKey,
+                            range: { min: minVal, max: maxVal },
+                            bucketWidth,
+                            skewness,
+                            kurtosis,
+                            histogram,
+                        });
+                    }
+                    // Build response with truncation indicators
+                    const response = {
+                        table: `${schema ?? "public"}.${table}`,
+                        column,
+                        groupBy,
+                        groups,
+                        count: groups.length,
+                    };
+                    // Add truncation indicators when groups are limited
+                    const groupsTruncated = effectiveGroupLimit !== undefined &&
+                        totalGroupCount > effectiveGroupLimit;
+                    if (groupsTruncated || !userProvidedGroupLimit) {
+                        response["truncated"] = groupsTruncated;
+                        response["totalGroupCount"] = totalGroupCount;
+                    }
+                    return response;
                 }
-                // Build response with truncation indicators
-                const response = {
+                // Ungrouped distribution (existing logic)
+                const moments = await computeMoments();
+                if (moments === null) {
+                    return { error: "No data or all nulls in column" };
+                }
+                const { minVal, maxVal, skewness, kurtosis } = moments;
+                const bucketWidth = Math.round(((maxVal - minVal) / numBuckets) * 1e6) / 1e6;
+                const histogram = await generateHistogram(minVal, maxVal);
+                return {
                     table: `${schema ?? "public"}.${table}`,
                     column,
-                    groupBy,
-                    groups,
-                    count: groups.length,
+                    range: { min: minVal, max: maxVal },
+                    bucketWidth,
+                    skewness,
+                    kurtosis,
+                    histogram,
                 };
-                // Add truncation indicators when groups are limited
-                const groupsTruncated = effectiveGroupLimit !== undefined &&
-                    totalGroupCount > effectiveGroupLimit;
-                if (groupsTruncated || !userProvidedGroupLimit) {
-                    response["truncated"] = groupsTruncated;
-                    response["totalGroupCount"] = totalGroupCount;
-                }
-                return response;
             }
-            // Ungrouped distribution (existing logic)
-            const moments = await computeMoments();
-            if (moments === null) {
-                return { error: "No data or all nulls in column" };
+            catch (error) {
+                return {
+                    success: false,
+                    error: formatPostgresError(error, { tool: "pg_stats_distribution" }),
+                };
             }
-            const { minVal, maxVal, skewness, kurtosis } = moments;
-            const bucketWidth = Math.round(((maxVal - minVal) / numBuckets) * 1e6) / 1e6;
-            const histogram = await generateHistogram(minVal, maxVal);
-            return {
-                table: `${schema ?? "public"}.${table}`,
-                column,
-                range: { min: minVal, max: maxVal },
-                bucketWidth,
-                skewness,
-                kurtosis,
-                histogram,
-            };
         },
     };
 }
@@ -650,87 +670,91 @@ export function createStatsHypothesisTool(adapter) {
         annotations: readOnly("Hypothesis Testing"),
         icons: getToolIcons("stats", readOnly("Hypothesis Testing")),
         handler: async (params, _context) => {
-            const { table, column, testType, hypothesizedMean, populationStdDev, schema, where, params: queryParams, groupBy, } = StatsHypothesisSchema.parse(params);
-            const schemaName = schema ?? "public";
-            const schemaPrefix = schema ? `"${schema}".` : "";
-            const whereClause = where ? `WHERE ${where}` : "";
-            // Validate column exists and is numeric
-            await validateNumericColumn(adapter, table, column, schemaName);
-            // Helper to calculate test results from row stats
-            const calculateTestResults = (n, sampleMean, sampleStdDev) => {
-                if (n < 2 || isNaN(sampleStdDev) || sampleStdDev === 0) {
-                    return { error: "Insufficient data or zero variance", sampleSize: n };
-                }
-                let stddevUsed;
-                let stddevNote;
-                if (testType === "z_test") {
-                    if (populationStdDev !== undefined) {
-                        stddevUsed = populationStdDev;
+            try {
+                const { table, column, testType, hypothesizedMean, populationStdDev, schema, where, params: queryParams, groupBy, } = StatsHypothesisSchema.parse(params);
+                const schemaName = schema ?? "public";
+                const schemaPrefix = schema ? `"${schema}".` : "";
+                const whereClause = where ? `WHERE ${where}` : "";
+                // Validate column exists and is numeric
+                await validateNumericColumn(adapter, table, column, schemaName);
+                // Helper to calculate test results from row stats
+                const calculateTestResults = (n, sampleMean, sampleStdDev) => {
+                    if (n < 2 || isNaN(sampleStdDev) || sampleStdDev === 0) {
+                        return {
+                            error: "Insufficient data or zero variance",
+                            sampleSize: n,
+                        };
+                    }
+                    let stddevUsed;
+                    let stddevNote;
+                    if (testType === "z_test") {
+                        if (populationStdDev !== undefined) {
+                            stddevUsed = populationStdDev;
+                        }
+                        else {
+                            stddevUsed = sampleStdDev;
+                            stddevNote =
+                                "No populationStdDev provided; using sample stddev (less accurate for z-test)";
+                        }
                     }
                     else {
                         stddevUsed = sampleStdDev;
-                        stddevNote =
-                            "No populationStdDev provided; using sample stddev (less accurate for z-test)";
                     }
-                }
-                else {
-                    stddevUsed = sampleStdDev;
-                }
-                const standardError = stddevUsed / Math.sqrt(n);
-                const testStatistic = (sampleMean - hypothesizedMean) / standardError;
-                const degreesOfFreedom = n - 1;
-                // Calculate p-value based on test type
-                const pValue = testType === "z_test"
-                    ? calculateZTestPValue(testStatistic)
-                    : calculateTTestPValue(testStatistic, degreesOfFreedom);
-                // Round p-value to 6 decimal places for cleaner output
-                const pValueRounded = Math.round(pValue * 1e6) / 1e6;
-                // Determine significance based on p-value
-                let interpretation;
-                if (pValueRounded < 0.001) {
-                    interpretation =
-                        "Highly significant (p < 0.001): Strong evidence against the null hypothesis";
-                }
-                else if (pValueRounded < 0.01) {
-                    interpretation =
-                        "Very significant (p < 0.01): Strong evidence against the null hypothesis";
-                }
-                else if (pValueRounded < 0.05) {
-                    interpretation =
-                        "Significant (p < 0.05): Evidence against the null hypothesis at α=0.05 level";
-                }
-                else if (pValueRounded < 0.1) {
-                    interpretation =
-                        "Marginally significant (p < 0.1): Weak evidence against the null hypothesis";
-                }
-                else {
-                    interpretation =
-                        "Not significant (p ≥ 0.1): Insufficient evidence to reject the null hypothesis";
-                }
-                // Build note with warnings
-                let noteText = stddevNote ??
-                    "Two-tailed p-value calculated using numerical approximation";
-                if (n < 30) {
-                    noteText =
-                        `Small sample size (n=${String(n)}): results may be less reliable. ` +
-                            noteText;
-                }
-                return {
-                    sampleSize: n,
-                    sampleMean,
-                    sampleStdDev,
-                    populationStdDev: testType === "z_test" ? (populationStdDev ?? null) : null,
-                    standardError,
-                    testStatistic,
-                    pValue: pValueRounded,
-                    degreesOfFreedom: testType === "t_test" ? degreesOfFreedom : null,
-                    interpretation,
-                    note: noteText,
+                    const standardError = stddevUsed / Math.sqrt(n);
+                    const testStatistic = (sampleMean - hypothesizedMean) / standardError;
+                    const degreesOfFreedom = n - 1;
+                    // Calculate p-value based on test type
+                    const pValue = testType === "z_test"
+                        ? calculateZTestPValue(testStatistic)
+                        : calculateTTestPValue(testStatistic, degreesOfFreedom);
+                    // Round p-value to 6 decimal places for cleaner output
+                    const pValueRounded = Math.round(pValue * 1e6) / 1e6;
+                    // Determine significance based on p-value
+                    let interpretation;
+                    if (pValueRounded < 0.001) {
+                        interpretation =
+                            "Highly significant (p < 0.001): Strong evidence against the null hypothesis";
+                    }
+                    else if (pValueRounded < 0.01) {
+                        interpretation =
+                            "Very significant (p < 0.01): Strong evidence against the null hypothesis";
+                    }
+                    else if (pValueRounded < 0.05) {
+                        interpretation =
+                            "Significant (p < 0.05): Evidence against the null hypothesis at α=0.05 level";
+                    }
+                    else if (pValueRounded < 0.1) {
+                        interpretation =
+                            "Marginally significant (p < 0.1): Weak evidence against the null hypothesis";
+                    }
+                    else {
+                        interpretation =
+                            "Not significant (p ≥ 0.1): Insufficient evidence to reject the null hypothesis";
+                    }
+                    // Build note with warnings
+                    let noteText = stddevNote ??
+                        "Two-tailed p-value calculated using numerical approximation";
+                    if (n < 30) {
+                        noteText =
+                            `Small sample size (n=${String(n)}): results may be less reliable. ` +
+                                noteText;
+                    }
+                    return {
+                        sampleSize: n,
+                        sampleMean,
+                        sampleStdDev,
+                        populationStdDev: testType === "z_test" ? (populationStdDev ?? null) : null,
+                        standardError,
+                        testStatistic,
+                        pValue: pValueRounded,
+                        degreesOfFreedom: testType === "t_test" ? degreesOfFreedom : null,
+                        interpretation,
+                        note: noteText,
+                    };
                 };
-            };
-            if (groupBy !== undefined) {
-                // Grouped hypothesis tests
-                const sql = `
+                if (groupBy !== undefined) {
+                    // Grouped hypothesis tests
+                    const sql = `
                     SELECT 
                         "${groupBy}" as group_key,
                         COUNT("${column}") as n,
@@ -741,59 +765,66 @@ export function createStatsHypothesisTool(adapter) {
                     GROUP BY "${groupBy}"
                     ORDER BY "${groupBy}"
                 `;
+                    const result = await adapter.executeQuery(sql, ...(queryParams !== undefined && queryParams.length > 0
+                        ? [queryParams]
+                        : []));
+                    const rows = result.rows ?? [];
+                    const groups = rows.map((row) => {
+                        const n = Number(row["n"]);
+                        const sampleMean = Number(row["mean"]);
+                        const sampleStdDev = Number(row["stddev"]);
+                        return {
+                            groupKey: row["group_key"],
+                            results: calculateTestResults(n, sampleMean, sampleStdDev),
+                        };
+                    });
+                    return {
+                        table: `${schema ?? "public"}.${table}`,
+                        column,
+                        testType,
+                        hypothesizedMean,
+                        groupBy,
+                        groups,
+                        count: groups.length,
+                    };
+                }
+                // Ungrouped hypothesis test
+                const sql = `
+                SELECT 
+                    COUNT("${column}") as n,
+                    AVG("${column}")::numeric(20,6) as mean,
+                    STDDEV_SAMP("${column}")::numeric(20,6) as stddev
+                FROM ${schemaPrefix}"${table}"
+                ${whereClause}
+            `;
                 const result = await adapter.executeQuery(sql, ...(queryParams !== undefined && queryParams.length > 0
                     ? [queryParams]
                     : []));
-                const rows = result.rows ?? [];
-                const groups = rows.map((row) => {
-                    const n = Number(row["n"]);
-                    const sampleMean = Number(row["mean"]);
-                    const sampleStdDev = Number(row["stddev"]);
-                    return {
-                        groupKey: row["group_key"],
-                        results: calculateTestResults(n, sampleMean, sampleStdDev),
-                    };
-                });
+                const row = result.rows?.[0];
+                if (!row)
+                    return { error: "No data found" };
+                const n = Number(row.n);
+                const sampleMean = Number(row.mean);
+                const sampleStdDev = Number(row.stddev);
+                const testResults = calculateTestResults(n, sampleMean, sampleStdDev);
+                // If error, return at top level (not nested in results)
+                if ("error" in testResults) {
+                    return testResults;
+                }
                 return {
                     table: `${schema ?? "public"}.${table}`,
                     column,
                     testType,
                     hypothesizedMean,
-                    groupBy,
-                    groups,
-                    count: groups.length,
+                    results: testResults,
                 };
             }
-            // Ungrouped hypothesis test
-            const sql = `
-                SELECT 
-                    COUNT("${column}") as n,
-                    AVG("${column}")::numeric(20,6) as mean,
-                    STDDEV_SAMP("${column}")::numeric(20,6) as stddev
-                FROM ${schemaPrefix}"${table}"
-                ${whereClause}
-            `;
-            const result = await adapter.executeQuery(sql, ...(queryParams !== undefined && queryParams.length > 0
-                ? [queryParams]
-                : []));
-            const row = result.rows?.[0];
-            if (!row)
-                return { error: "No data found" };
-            const n = Number(row.n);
-            const sampleMean = Number(row.mean);
-            const sampleStdDev = Number(row.stddev);
-            const testResults = calculateTestResults(n, sampleMean, sampleStdDev);
-            // If error, return at top level (not nested in results)
-            if ("error" in testResults) {
-                return testResults;
+            catch (error) {
+                return {
+                    success: false,
+                    error: formatPostgresError(error, { tool: "pg_stats_hypothesis" }),
+                };
             }
-            return {
-                table: `${schema ?? "public"}.${table}`,
-                column,
-                testType,
-                hypothesizedMean,
-                results: testResults,
-            };
         },
     };
 }
@@ -810,105 +841,113 @@ export function createStatsSamplingTool(adapter) {
         annotations: readOnly("Random Sampling"),
         icons: getToolIcons("stats", readOnly("Random Sampling")),
         handler: async (params, _context) => {
-            const { table, method, sampleSize, percentage, schema, select, where, params: queryParams, } = StatsSamplingSchema.parse(params);
-            const schemaName = schema ?? "public";
-            // Validate table exists
-            await validateTableExists(adapter, table, schemaName);
-            const schemaPrefix = schema ? `"${schema}".` : "";
-            const columns = select && select.length > 0
-                ? select.map((c) => `"${c}"`).join(", ")
-                : "*";
-            const whereClause = where ? `WHERE ${where}` : "";
-            const samplingMethod = method ?? "random";
-            let sql;
-            let note;
-            // If sampleSize is provided, always use ORDER BY RANDOM() LIMIT n for exact counts
-            // TABLESAMPLE BERNOULLI/SYSTEM are percentage-based and cannot guarantee exact row counts
-            if (sampleSize !== undefined) {
-                const limit = sampleSize;
-                sql = `
+            try {
+                const { table, method, sampleSize, percentage, schema, select, where, params: queryParams, } = StatsSamplingSchema.parse(params);
+                const schemaName = schema ?? "public";
+                // Validate table exists
+                await validateTableExists(adapter, table, schemaName);
+                const schemaPrefix = schema ? `"${schema}".` : "";
+                const columns = select && select.length > 0
+                    ? select.map((c) => `"${c}"`).join(", ")
+                    : "*";
+                const whereClause = where ? `WHERE ${where}` : "";
+                const samplingMethod = method ?? "random";
+                let sql;
+                let note;
+                // If sampleSize is provided, always use ORDER BY RANDOM() LIMIT n for exact counts
+                // TABLESAMPLE BERNOULLI/SYSTEM are percentage-based and cannot guarantee exact row counts
+                if (sampleSize !== undefined) {
+                    const limit = sampleSize;
+                    sql = `
                     SELECT ${columns}
                     FROM ${schemaPrefix}"${table}"
                     ${whereClause}
                     ORDER BY RANDOM()
                     LIMIT ${String(limit)}
                 `;
-                if (percentage !== undefined) {
-                    note = `sampleSize (${String(sampleSize)}) takes precedence over percentage (${String(percentage)}%). Using ORDER BY RANDOM() LIMIT for exact row count.`;
-                }
-                else if (samplingMethod !== "random") {
-                    note = `Using ORDER BY RANDOM() LIMIT for exact ${String(sampleSize)} row count. TABLESAMPLE ${samplingMethod.toUpperCase()} is percentage-based and cannot guarantee exact counts.`;
+                    if (percentage !== undefined) {
+                        note = `sampleSize (${String(sampleSize)}) takes precedence over percentage (${String(percentage)}%). Using ORDER BY RANDOM() LIMIT for exact row count.`;
+                    }
+                    else if (samplingMethod !== "random") {
+                        note = `Using ORDER BY RANDOM() LIMIT for exact ${String(sampleSize)} row count. TABLESAMPLE ${samplingMethod.toUpperCase()} is percentage-based and cannot guarantee exact counts.`;
+                    }
                 }
-            }
-            else if (samplingMethod === "random") {
-                // Default random sampling with default limit (20 to reduce LLM context usage)
-                const limit = 20;
-                sql = `
+                else if (samplingMethod === "random") {
+                    // Default random sampling with default limit (20 to reduce LLM context usage)
+                    const limit = 20;
+                    sql = `
                     SELECT ${columns}
                     FROM ${schemaPrefix}"${table}"
                     ${whereClause}
                     ORDER BY RANDOM()
                     LIMIT ${String(limit)}
                 `;
-                if (percentage !== undefined) {
-                    note = `percentage (${String(percentage)}%) is ignored for random method. Use method:'bernoulli' or method:'system' for percentage-based sampling, or use sampleSize for exact row count.`;
+                    if (percentage !== undefined) {
+                        note = `percentage (${String(percentage)}%) is ignored for random method. Use method:'bernoulli' or method:'system' for percentage-based sampling, or use sampleSize for exact row count.`;
+                    }
                 }
-            }
-            else {
-                // TABLESAMPLE with percentage (approximate row count)
-                // Apply default limit to prevent large payloads
-                const pct = percentage ?? 10;
-                const DEFAULT_TABLESAMPLE_LIMIT = 100;
-                sql = `
+                else {
+                    // TABLESAMPLE with percentage (approximate row count)
+                    // Apply default limit to prevent large payloads
+                    const pct = percentage ?? 10;
+                    const DEFAULT_TABLESAMPLE_LIMIT = 100;
+                    sql = `
                     SELECT ${columns}
                     FROM ${schemaPrefix}"${table}"
                     TABLESAMPLE ${samplingMethod.toUpperCase()}(${String(pct)})
                     ${whereClause}
                     LIMIT ${String(DEFAULT_TABLESAMPLE_LIMIT + 1)}
                 `;
-                // Add hint about system method unreliability for small tables
-                const methodHint = samplingMethod === "system"
-                    ? " Consider using 'bernoulli' or 'random' method for more reliable results on small tables."
-                    : "";
-                note = `TABLESAMPLE ${samplingMethod.toUpperCase()}(${String(pct)}%) returns approximately ${String(pct)}% of rows. Actual count varies based on table size and sampling algorithm.${methodHint}`;
-            }
-            const result = await adapter.executeQuery(sql, ...(queryParams !== undefined && queryParams.length > 0
-                ? [queryParams]
-                : []));
-            let rows = result.rows ?? [];
-            // Check if we need to truncate due to default limit for TABLESAMPLE methods
-            let truncated = false;
-            let totalSampled;
-            const DEFAULT_TABLESAMPLE_LIMIT = 100;
-            if (sampleSize === undefined &&
-                samplingMethod !== "random" &&
-                rows.length > DEFAULT_TABLESAMPLE_LIMIT) {
-                totalSampled = rows.length;
-                rows = rows.slice(0, DEFAULT_TABLESAMPLE_LIMIT);
-                truncated = true;
-            }
-            const response = {
-                table: `${schema ?? "public"}.${table}`,
-                method: samplingMethod,
-                sampleSize: rows.length,
-                rows,
-            };
-            // Add truncation indicators if applicable
-            if (truncated && totalSampled !== undefined) {
-                response.truncated = truncated;
-                response.totalSampled = totalSampled;
-            }
-            if (note !== undefined) {
-                response.note = note;
+                    // Add hint about system method unreliability for small tables
+                    const methodHint = samplingMethod === "system"
+                        ? " Consider using 'bernoulli' or 'random' method for more reliable results on small tables."
+                        : "";
+                    note = `TABLESAMPLE ${samplingMethod.toUpperCase()}(${String(pct)}%) returns approximately ${String(pct)}% of rows. Actual count varies based on table size and sampling algorithm.${methodHint}`;
+                }
+                const result = await adapter.executeQuery(sql, ...(queryParams !== undefined && queryParams.length > 0
+                    ? [queryParams]
+                    : []));
+                let rows = result.rows ?? [];
+                // Check if we need to truncate due to default limit for TABLESAMPLE methods
+                let truncated = false;
+                let totalSampled;
+                const DEFAULT_TABLESAMPLE_LIMIT = 100;
+                if (sampleSize === undefined &&
+                    samplingMethod !== "random" &&
+                    rows.length > DEFAULT_TABLESAMPLE_LIMIT) {
+                    totalSampled = rows.length;
+                    rows = rows.slice(0, DEFAULT_TABLESAMPLE_LIMIT);
+                    truncated = true;
+                }
+                const response = {
+                    table: `${schema ?? "public"}.${table}`,
+                    method: samplingMethod,
+                    sampleSize: rows.length,
+                    rows,
+                };
+                // Add truncation indicators if applicable
+                if (truncated && totalSampled !== undefined) {
+                    response.truncated = truncated;
+                    response.totalSampled = totalSampled;
+                }
+                if (note !== undefined) {
+                    response.note = note;
+                }
+                // Add note if requested sampleSize exceeded available rows
+                if (sampleSize !== undefined && rows.length < sampleSize) {
+                    const existingNote = response.note !== undefined ? response.note + " " : "";
+                    response.note =
+                        existingNote +
+                            `Requested ${String(sampleSize)} rows but only ${String(rows.length)} available.`;
+                }
+                return response;
             }
-            // Add note if requested sampleSize exceeded available rows
-            if (sampleSize !== undefined && rows.length < sampleSize) {
-                const existingNote = response.note !== undefined ? response.note + " " : "";
-                response.note =
-                    existingNote +
-                        `Requested ${String(sampleSize)} rows but only ${String(rows.length)} available.`;
+            catch (error) {
+                return {
+                    success: false,
+                    error: formatPostgresError(error, { tool: "pg_stats_sampling" }),
+                };
             }
-            return response;
         },
     };
 }