pg_query 6.1.0 → 6.2.1

data/ext/pg_query/pg_query_summary_truncate.c

@@ -0,0 +1,530 @@
+#include "pg_query.h"
+#include "pg_query_internal.h"
+#include "pg_query_summary.h"
+
+#include "nodes/pg_list.h"
+#include "nodes/nodeFuncs.h"
+
+#include "mb/pg_wchar.h"
+
+enum TruncationAttr
+{
+	TRUNCATION_TARGET_LIST,
+	TRUNCATION_WHERE_CLAUSE,
+	TRUNCATION_VALUES_LISTS,
+	TRUNCATION_COLS,
+	TRUNCATION_CTE_QUERY,
+};
+
+typedef struct
+{
+	List	   *truncations;
+	int32_t		depth;
+}			TruncationState;
+
+typedef struct
+{
+	enum TruncationAttr attr;
+	Node	   *node;
+	int32_t		depth;
+	int32_t		length;
+}			PossibleTruncation;
+
+static bool generate_possible_truncations(Node *tree, TruncationState * state);
+static void apply_truncations(Summary * summary, Node *tree, TruncationState * state, int truncate_limit);
+static int	cmp_possible_truncations(const ListCell *a, const ListCell *b);
+
+static int32_t select_target_list_len(List *nodes);
+static int32_t select_values_lists_len(List *nodes);
+static int32_t update_target_list_len(List *nodes);
+static int32_t where_clause_len(Node *node);
+static int32_t cols_len(List *nodes);
+
+static ColumnRef *dummy_column(void);
+static ResTarget *dummy_target(void);
+static Node *dummy_select(List *targetList, Node *whereClause, List *valuesLists);
+static Node *dummy_insert(List *cols);
+static Node *dummy_update(List *targetList);
+
+static size_t deparse_stmt_len(Node *node);
+static StringInfo deparse_raw_stmt_list(List *stmts);
+
+/*  Given a walked parse tree and summary, store the truncated version in `summary`. */
+void
+pg_query_summary_truncate(Summary * summary, Node *tree, int truncate_limit)
+{
+	TruncationState state = {NULL, 0};
+
+	StringInfo	output = deparse_raw_stmt_list((List *) tree);
+
+	if (output->len <= truncate_limit)
+	{
+		summary->truncated_query = output->data;
+		return;
+	}
+
+	destroyStringInfo(output);
+
+	generate_possible_truncations(tree, &state);
+
+	list_sort(state.truncations, cmp_possible_truncations);
+	apply_truncations(summary, tree, &state, truncate_limit);
+}
+
+static void
+truncate_mbstr(char *mbstr, size_t max_chars)
+{
+	/* Determine the number of characters in mbstr. */
+	int			n_chars = pg_mbstrlen(mbstr);
+
+	/* If we don't need to truncate the string, return immediately. */
+	if (n_chars <= max_chars)
+		return;
+
+	/* Determine how many bytes hold `max_chars - 3`. */
+	int			n_bytes = pg_mbcharcliplen(mbstr, n_chars, max_chars - 3);
+
+	/* Actually truncate it. */
+	strncpy(mbstr + n_bytes, "...", 4);
+	mbstr[n_bytes + 3] = '\0';
+}
+
+static void
+add_truncation(TruncationState * state, enum TruncationAttr attr,
+			   Node *node, int32_t length)
+{
+	/* Don't bother truncating if it won't become shorter. */
+	if (length <= 3)
+		return;
+
+	PossibleTruncation *truncation = palloc(sizeof(PossibleTruncation));
+
+	truncation->attr = attr;
+	truncation->node = node;
+	truncation->depth = state->depth;
+	truncation->length = length;
+
+	state->truncations = lappend(state->truncations, truncation);
+}
+
+static void
+add_truncation_where_clause(TruncationState * state, Node *node, Node *whereClause)
+{
+	if (whereClause == NULL)
+		return;
+
+	add_truncation(state,
+				   TRUNCATION_WHERE_CLAUSE,
+				   node,
+				   where_clause_len(whereClause));
+}
+
+static bool
+generate_possible_truncations(Node *node, TruncationState * state)
+{
+	if (node == NULL)
+		return false;
+
+	switch (nodeTag(node))
+	{
+		case T_RawStmt:
+			return generate_possible_truncations(castNode(RawStmt, node)->stmt, state);
+
+		case T_SelectStmt:
+			{
+				SelectStmt *stmt = castNode(SelectStmt, node);
+
+				if (stmt->targetList != NULL)
+					add_truncation(state,
+								   TRUNCATION_TARGET_LIST,
+								   node,
+								   select_target_list_len(stmt->targetList));
+
+				add_truncation_where_clause(state, node, stmt->whereClause);
+
+				if (stmt->valuesLists != NULL)
+					add_truncation(state,
+								   TRUNCATION_VALUES_LISTS,
+								   node,
+								   select_values_lists_len(stmt->valuesLists));
+
+				break;
+			}
+
+		case T_InsertStmt:
+			{
+				InsertStmt *stmt = castNode(InsertStmt, node);
+
+				if (stmt->cols != NULL)
+					add_truncation(state, TRUNCATION_COLS, node, cols_len(stmt->cols));
+
+				break;
+			}
+
+		case T_UpdateStmt:
+			{
+				UpdateStmt *stmt = castNode(UpdateStmt, node);
+
+				if (stmt->targetList != NULL)
+					add_truncation(state,
+								   TRUNCATION_TARGET_LIST,
+								   node,
+								   update_target_list_len(stmt->targetList));
+
+				add_truncation_where_clause(state, node, stmt->whereClause);
+
+				break;
+			}
+
+		case T_DeleteStmt:
+			{
+				DeleteStmt *stmt = castNode(DeleteStmt, node);
+
+				add_truncation_where_clause(state, node, stmt->whereClause);
+				break;
+			}
+
+		case T_CopyStmt:
+			{
+				CopyStmt   *stmt = castNode(CopyStmt, node);
+
+				add_truncation_where_clause(state, node, stmt->whereClause);
+				break;
+			}
+
+		case T_IndexStmt:
+			{
+				IndexStmt  *stmt = castNode(IndexStmt, node);
+
+				add_truncation_where_clause(state, node, stmt->whereClause);
+				break;
+			}
+
+		case T_RuleStmt:
+			{
+				RuleStmt   *stmt = castNode(RuleStmt, node);
+
+				add_truncation_where_clause(state, node, stmt->whereClause);
+				break;
+			}
+
+		case T_CommonTableExpr:
+			{
+				CommonTableExpr *stmt = castNode(CommonTableExpr, node);
+
+				if (stmt->ctequery != NULL)
+				{
+					size_t		query_len = deparse_stmt_len((Node *) stmt->ctequery);
+
+					add_truncation(state,
+								   TRUNCATION_CTE_QUERY,
+								   node,
+								   query_len);
+				}
+
+				break;
+			}
+
+		case T_InferClause:
+			{
+				InferClause *stmt = castNode(InferClause, node);
+
+				add_truncation_where_clause(state, node, stmt->whereClause);
+				break;
+			}
+
+		case T_OnConflictClause:
+			{
+				OnConflictClause *stmt = castNode(OnConflictClause, node);
+
+				if (stmt->targetList != NULL)
+					add_truncation(state,
+								   TRUNCATION_TARGET_LIST,
+								   node,
+								   update_target_list_len(stmt->targetList));
+
+				add_truncation_where_clause(state, node, stmt->whereClause);
+
+				break;
+			}
+
+		default:
+			break;
+	}
+
+	if (!pg_query_raw_tree_walker_supports(node))
+		return false;
+
+	int			old_depth = state->depth;
+
+	state->depth++;
+
+	bool		result = raw_expression_tree_walker(node, generate_possible_truncations, (void *) state);
+
+	/* Restore old depth value, since the current node (or its parents) may */
+	/* have sibling elements. */
+	state->depth = old_depth;
+
+	return result;
+}
+
+static int
+cmp_possible_truncations(const ListCell *a, const ListCell *b)
+{
+	const		PossibleTruncation *pt_a = lfirst(a);
+	const		PossibleTruncation *pt_b = lfirst(b);
+
+	int			depth_cmp = pt_b->depth - pt_a->depth;
+
+	if (depth_cmp != 0)
+		return depth_cmp;
+
+	return pt_b->length - pt_a->length;
+}
+
+static void
+global_replace(char *str, char *pattern, char *replacement)
+{
+	size_t		plen = strlen(pattern);
+	size_t		rlen = strlen(replacement);
+	char	   *s = str;
+
+	Assert(plen >= rlen);
+
+	while ((s = strstr(s, pattern)) != NULL)
+	{
+		memcpy(s, replacement, rlen);
+
+		/* Shift remainder of the string if needed */
+		if (plen > rlen)
+			memmove(s + rlen, s + plen, strlen(s + plen) + 1);
+
+		s += rlen;
+	}
+}
+
+static void
+apply_truncations(Summary * summary, Node *tree, TruncationState * state, int truncation_limit)
+{
+	List	   *truncations = state->truncations;
+	StringInfo	output = NULL;
+	ListCell   *lc;
+
+	foreach(lc, state->truncations)
+	{
+		PossibleTruncation *truncation = lfirst(lc);
+
+		Node	   *node = truncation->node;
+		enum TruncationAttr attr = truncation->attr;
+
+		if (IsA(node, SelectStmt) && attr == TRUNCATION_TARGET_LIST)
+			castNode(SelectStmt, node)->targetList = list_make1(dummy_target());
+		else if (IsA(node, SelectStmt) && attr == TRUNCATION_WHERE_CLAUSE)
+			castNode(SelectStmt, node)->whereClause = (Node *) dummy_column();
+		else if (IsA(node, SelectStmt) && attr == TRUNCATION_VALUES_LISTS)
+			castNode(SelectStmt, node)->valuesLists = list_make1(list_make1(dummy_column()));
+		else if (IsA(node, UpdateStmt) && attr == TRUNCATION_TARGET_LIST)
+			castNode(UpdateStmt, node)->targetList = list_make1(dummy_target());
+		else if (IsA(node, InsertStmt) && attr == TRUNCATION_COLS)
+			castNode(InsertStmt, node)->cols = list_make1(dummy_target());
+		else if (IsA(node, UpdateStmt) && attr == TRUNCATION_WHERE_CLAUSE)
+			castNode(UpdateStmt, node)->whereClause = (Node *) dummy_column();
+		else if (IsA(node, DeleteStmt) && attr == TRUNCATION_WHERE_CLAUSE)
+			castNode(DeleteStmt, node)->whereClause = (Node *) dummy_column();
+		else if (IsA(node, CopyStmt) && attr == TRUNCATION_WHERE_CLAUSE)
+			castNode(CopyStmt, node)->whereClause = (Node *) dummy_column();
+		else if (IsA(node, IndexStmt) && attr == TRUNCATION_WHERE_CLAUSE)
+			castNode(IndexStmt, node)->whereClause = (Node *) dummy_column();
+		else if (IsA(node, RuleStmt) && attr == TRUNCATION_WHERE_CLAUSE)
+			castNode(RuleStmt, node)->whereClause = (Node *) dummy_column();
+		else if (IsA(node, CommonTableExpr) && attr == TRUNCATION_CTE_QUERY)
+			castNode(CommonTableExpr, node)->ctequery = dummy_select(NULL, (Node *) dummy_column(), NULL);
+		else if (IsA(node, InferClause) && attr == TRUNCATION_WHERE_CLAUSE)
+			castNode(InferClause, node)->whereClause = (Node *) dummy_column();
+		else if (IsA(node, OnConflictClause) && attr == TRUNCATION_TARGET_LIST)
+			castNode(OnConflictClause, node)->targetList = list_make1(dummy_target());
+		else if (IsA(node, OnConflictClause) && attr == TRUNCATION_WHERE_CLAUSE)
+			castNode(OnConflictClause, node)->whereClause = (Node *) dummy_column();
+		else
+			elog(ERROR, "apply_truncations() got unknown truncation type");
+
+		if (output)
+			destroyStringInfo(output);
+		output = deparse_raw_stmt_list((List *) tree);
+
+		global_replace(output->data, "SELECT \"…\" AS \"…\"", "SELECT \"…\"");
+		global_replace(output->data, "SELECT WHERE \"…\"", "\"…\"");
+		global_replace(output->data, "\"…\"", "...");
+
+		if (strlen(output->data) <= truncation_limit)
+		{
+			summary->truncated_query = output->data;
+			return;
+		}
+	}
+
+	if (!output)
+		output = deparse_raw_stmt_list((List *) tree);
+
+	truncate_mbstr(output->data, truncation_limit);
+	summary->truncated_query = output->data;
+}
+
+static ColumnRef *
+dummy_column(void)
+{
+	ColumnRef  *colref = makeNode(ColumnRef);
+
+	colref->fields = list_make1(makeString(pstrdup("…")));
+	colref->location = 0;
+
+	return colref;
+}
+
+static ResTarget *
+dummy_target(void)
+{
+	ResTarget  *target = makeNode(ResTarget);
+
+	target->name = pstrdup("…");
+	target->val = (Node *) dummy_column();
+
+	/*
+	 * TODO:docs for ResTarget say "-1 if unknown"-- would that be more
+	 * correct ? (see also, dummy_column())
+	 */
+	target->location = 0;
+
+	return target;
+}
+
+static Node *
+dummy_select(List *targetList, Node *whereClause, List *valuesLists)
+{
+	SelectStmt *stmt = makeNode(SelectStmt);
+
+	stmt->targetList = targetList;
+	stmt->whereClause = whereClause;
+	stmt->valuesLists = valuesLists;
+
+	return (Node *) stmt;
+}
+
+static Node *
+dummy_insert(List *cols)
+{
+	RangeVar   *rv = makeNode(RangeVar);
+
+	rv->relname = pstrdup("x");
+	rv->inh = true;
+	rv->relpersistence = 'p';
+	rv->location = 0;
+
+	InsertStmt *stmt = makeNode(InsertStmt);
+
+	stmt->relation = rv;
+	stmt->cols = cols;
+	stmt->override = 1;
+
+	return (Node *) stmt;
+}
+
+static Node *
+dummy_update(List *targetList)
+{
+	RangeVar   *rv = makeNode(RangeVar);
+
+	rv->relname = pstrdup("x");
+	rv->inh = true;
+	rv->relpersistence = 'p';
+	rv->location = 0;
+
+	UpdateStmt *stmt = makeNode(UpdateStmt);
+
+	stmt->relation = rv;
+	stmt->targetList = targetList;
+
+	return (Node *) stmt;
+}
+
+static int32_t
+select_target_list_len(List *nodes)
+{
+	size_t		result_len = deparse_stmt_len(dummy_select(nodes, NULL, NULL));
+	size_t		dummy_len = 7;	/* "SELECT " */
+
+	return (int32_t) result_len - dummy_len;
+}
+
+static int32_t
+select_values_lists_len(List *nodes)
+{
+	size_t		result_len = deparse_stmt_len(dummy_select(NULL, NULL, nodes));
+	size_t		dummy_len = 9;	/* "VALUES ()" */
+
+	return (int32_t) result_len - dummy_len;
+}
+
+static int32_t
+update_target_list_len(List *nodes)
+{
+	size_t		result_len = deparse_stmt_len(dummy_update(nodes));
+	size_t		dummy_len = 13; /* "UPDATE x SET " */
+
+	return (int32_t) result_len - dummy_len;
+}
+
+static int32_t
+where_clause_len(Node *node)
+{
+	size_t		result_len = deparse_stmt_len(dummy_select(NULL, node, NULL));
+	size_t		dummy_len = 13; /* "SELECT WHERE " */
+
+	return (int32_t) result_len - dummy_len;
+}
+
+static int32_t
+cols_len(List *nodes)
+{
+	size_t		result_len = deparse_stmt_len(dummy_insert(nodes));
+	size_t		dummy_len = 31; /* "INSERT INTO x () DEFAULT VALUES" */
+
+	return (int32_t) result_len - dummy_len;
+}
+
+static size_t
+deparse_stmt_len(Node *node)
+{
+	if (!IsA(node, RawStmt))
+	{
+		RawStmt    *raw = makeNode(RawStmt);
+
+		raw->stmt = node;
+		raw->stmt_location = -1;
+		raw->stmt_len = 0;
+		node = (Node *) raw;
+	}
+
+	StringInfo	str = deparse_raw_stmt_list(list_make1(node));
+	size_t		len = str->len;
+
+	destroyStringInfo(str);
+	return len;
+}
+
+static StringInfo
+deparse_raw_stmt_list(List *stmts)
+{
+	PostgresDeparseOpts opts = {0};
+	StringInfo	str = makeStringInfo();
+	ListCell   *lc;
+
+	foreach(lc, stmts)
+	{
+		deparseRawStmtOpts(str, castNode(RawStmt, lfirst(lc)), opts);
+		if (lnext(stmts, lc))
+			appendStringInfoString(str, "; ");
+	}
+
+	return str;
+}