Add a callback on an action of a transaction cleaning.

Now AQO could clean global data, such of query_text at the end
of transaction. It is needed because errors during query execution
phases. We need to be sure that all such data cleaned before the next
query execution.
Interface of the add_query_text routine changed.

Author: danolivo

aqo.c

@@ -111,6 +111,22 @@ ExplainOneNode_hook_type					prev_ExplainOneNode_hook;
  *
  *****************************************************************************/
 
+static void
+aqo_free_callback(ResourceReleasePhase phase,
+					 bool isCommit,
+					 bool isTopLevel,
+					 void *arg)
+{
+	if (phase != RESOURCE_RELEASE_AFTER_LOCKS)
+		return;
+
+	if (query_text != NULL)
+	{
+		pfree(query_text);
+		query_text = NULL;
+	}
+}
+
 void
 _PG_init(void)
 {
@@ -220,6 +236,7 @@ _PG_init(void)
 	AQOMemoryContext = AllocSetContextCreate(TopMemoryContext,
 											 "AQOMemoryContext",
 											 ALLOCSET_DEFAULT_SIZES);
+	RegisterResourceReleaseCallback(aqo_free_callback, NULL);
 }
 
 PG_FUNCTION_INFO_V1(invalidate_deactivated_queries_cache);

aqo.h

@@ -286,7 +286,7 @@ int			get_clause_hash(Expr *clause, int nargs,
 extern bool find_query(int qhash, Datum *search_values, bool *search_nulls);
 extern bool update_query(int qhash, int fhash,
 						 bool learn_aqo, bool use_aqo, bool auto_tuning);
-extern bool add_query_text(int query_hash, char *query_text);
+extern bool add_query_text(int query_hash);
 extern bool load_fss(int fhash, int fss_hash,
 					 int ncols, double **matrix, double *targets, int *rows);
 extern bool update_fss(int fhash, int fss_hash, int nrows, int ncols,

preprocessing.c

@@ -71,16 +71,19 @@ static bool isQueryUsingSystemRelation_walker(Node *node, void *context);
 void
 get_query_text(ParseState *pstate, Query *query)
 {
-	MemoryContext	oldCxt;
-
 	/*
 	 * Duplicate query string into private AQO memory context for guard
 	 * from possible memory context switching.
 	 */
-	oldCxt = MemoryContextSwitchTo(AQOMemoryContext);
 	if (pstate)
+	{
+		MemoryContext oldCxt = MemoryContextSwitchTo(AQOMemoryContext);
 		query_text = pstrdup(pstate->p_sourcetext);
-	MemoryContextSwitchTo(oldCxt);
+		MemoryContextSwitchTo(oldCxt);
+	}
+	else
+		/* Can't imagine such case. Still, throw an error. */
+		elog(ERROR, "[AQO]: Query text is not found in post-parse step");
 
 	if (prev_post_parse_analyze_hook)
 		prev_post_parse_analyze_hook(pstate, query);
@@ -97,7 +100,7 @@ call_default_planner(Query *parse, int cursorOptions, ParamListInfo boundParams)
 	else
 		return standard_planner(parse, cursorOptions, boundParams);
 }
-#include "tcop/tcopprot.h"
+
 /*
  * Before query optimization we determine machine learning settings
  * for the query.
@@ -218,8 +221,12 @@ aqo_planner(Query *parse, int cursorOptions, ParamListInfo boundParams)
 							  query_context.use_aqo,
 							  query_context.auto_tuning);
 
-
-			add_query_text(query_context.query_hash, query_text);
+			/*
+			 * Add query text into the ML-knowledge base. Just for further
+			 * analysis. In the case of cached plans we could have NULL query text.
+			 */
+			if (query_text != NULL)
+				add_query_text(query_context.query_hash);
 		}
 	}
 	else

storage.c

@@ -214,7 +214,7 @@ update_query(int qhash, int fhash,
  * Returns false if the operation failed, true otherwise.
  */
 bool
-add_query_text(int qhash, char *query_text)
+add_query_text(int qhash)
 {
 	RangeVar   *rv;
 	Relation	hrel;