通常のクエリでもplpgsqlのクエリのように扱ってしまうバグを修正した。

[pghintplan/pg_hint_plan.git] / pg_hint_plan.c

diff --git a/pg_hint_plan.c b/pg_hint_plan.c
index 51b58f0..db1082a 100644 (file)
--- a/pg_hint_plan.c
+++ b/pg_hint_plan.c
@@ -36,6 +36,8 @@
 #include "catalog/pg_class.h"
 #endif
 
+#include "plpgsql.h"
+
 #ifdef PG_MODULE_MAGIC
 PG_MODULE_MAGIC;
 #endif
@@ -384,6 +386,13 @@ static void set_dummy_rel_pathlist(RelOptInfo *rel);
 RelOptInfo *pg_hint_plan_make_join_rel(PlannerInfo *root, RelOptInfo *rel1,
                                                                           RelOptInfo *rel2);
 
+static void pg_hint_plan_plpgsql_func_setup(PLpgSQL_execstate *estate,
+                                                                                 PLpgSQL_stmt *stmt);
+static void pg_hint_plan_plpgsql_stmt_beg(PLpgSQL_execstate *estate,
+                                                                                 PLpgSQL_stmt *stmt);
+static void pg_hint_plan_plpgsql_stmt_end(PLpgSQL_execstate *estate,
+                                                                                 PLpgSQL_stmt *stmt);
+
 /* GUC variables */
 static bool    pg_hint_plan_enable_hint = true;
 static bool    pg_hint_plan_debug_print = false;
@@ -458,6 +467,8 @@ static const HintParser parsers[] = {
        {NULL, NULL, HINT_KEYWORD_UNRECOGNIZED}
 };
 
+const char *hint_query_string = NULL;
+PLpgSQL_plugin  plugin_funcs = { };
 /*
  * Module load callbacks
  */
@@ -508,6 +519,11 @@ _PG_init(void)
        get_relation_info_hook = pg_hint_plan_get_relation_info;
        prev_join_search = join_search_hook;
        join_search_hook = pg_hint_plan_join_search;
+
+       /* PL/pgSQL plugin hook */
+       PLpgSQL_plugin  **var_ptr = (PLpgSQL_plugin **) find_rendezvous_variable("PLpgSQL_plugin");
+       *var_ptr = &plugin_funcs;
+       (&plugin_funcs)->func_setup = (void *)pg_hint_plan_plpgsql_func_setup;
 }
 
 /*
@@ -1342,6 +1358,8 @@ parse_head_comment(Query *parse)
                entry = FetchPreparedStatement(stmt_name, true);
                p = entry->plansource->query_string;
        }
+       else if (hint_query_string)
+               p = hint_query_string;
        else
                p = debug_query_string;
 
@@ -3352,6 +3370,30 @@ set_rel_pathlist(PlannerInfo *root, RelOptInfo *rel,
        }
 }
 
+static void
+pg_hint_plan_plpgsql_func_setup(PLpgSQL_execstate *estate, PLpgSQL_stmt *stmt)
+{
+       (&plugin_funcs)->stmt_beg = pg_hint_plan_plpgsql_stmt_beg;
+       (&plugin_funcs)->stmt_end = pg_hint_plan_plpgsql_stmt_end;
+}
+
+static void
+pg_hint_plan_plpgsql_stmt_beg(PLpgSQL_execstate *estate, PLpgSQL_stmt *stmt)
+{
+       if ((enum PLpgSQL_stmt_types) stmt->cmd_type == PLPGSQL_STMT_EXECSQL)
+       {
+               PLpgSQL_expr *expr = ((PLpgSQL_stmt_execsql *) stmt)->sqlstmt;
+               hint_query_string = expr->query;
+       }
+}
+
+static void
+pg_hint_plan_plpgsql_stmt_end(PLpgSQL_execstate *estate, PLpgSQL_stmt *stmt)
+{
+       if ((enum PLpgSQL_stmt_types) stmt->cmd_type == PLPGSQL_STMT_EXECSQL)
+               hint_query_string = NULL;
+}
+
 #define standard_join_search pg_hint_plan_standard_join_search
 #define join_search_one_level pg_hint_plan_join_search_one_level
 #define make_join_rel make_join_rel_wrapper