44
55#include "postgres_deparse.h"
66
7+ #include "postgres.h"
8+ #include "lib/stringinfo.h"
9+ #include "nodes/parsenodes.h"
10+
11+ #include "protobuf/pg_query.pb-c.h"
12+
713PgQueryDeparseResult pg_query_deparse_protobuf (PgQueryProtobuf parse_tree )
14+ {
15+ PostgresDeparseOpts opts ;
16+ MemSet (& opts , 0 , sizeof (PostgresDeparseOpts ));
17+ return pg_query_deparse_protobuf_opts (parse_tree , opts );
18+ }
19+
20+ PgQueryDeparseResult pg_query_deparse_protobuf_opts (PgQueryProtobuf parse_tree , PostgresDeparseOpts opts )
821{
922 PgQueryDeparseResult result = {0 };
1023 StringInfoData str ;
@@ -21,7 +34,7 @@ PgQueryDeparseResult pg_query_deparse_protobuf(PgQueryProtobuf parse_tree)
2134 initStringInfo (& str );
2235
2336 foreach (lc , stmts ) {
24- deparseRawStmt (& str , castNode (RawStmt , lfirst (lc )));
37+ deparseRawStmtOpts (& str , castNode (RawStmt , lfirst (lc )), opts );
2538 if (lnext (stmts , lc ))
2639 appendStringInfoString (& str , "; " );
2740 }
@@ -62,3 +75,103 @@ void pg_query_free_deparse_result(PgQueryDeparseResult result)
6275
6376 free (result .query );
6477}
78+
79+ PgQueryDeparseCommentsResult pg_query_deparse_comments_for_query (const char * query )
80+ {
81+ PgQueryDeparseCommentsResult result = {0 };
82+ PgQueryScanResult scan_result_raw = pg_query_scan (query );
83+ if (scan_result_raw .error )
84+ {
85+ result .error = scan_result_raw .error ;
86+ return result ;
87+ }
88+
89+ PgQuery__ScanResult * scan_result = pg_query__scan_result__unpack (NULL , scan_result_raw .pbuf .len , (void * ) scan_result_raw .pbuf .data );
90+ bool prior_token_was_comment = false;
91+ int32_t prior_non_comment_end = 0 ;
92+ int32_t prior_token_end = 0 ;
93+
94+ result .comment_count = 0 ;
95+ for (int i = 0 ; i < scan_result -> n_tokens ; i ++ )
96+ {
97+ PgQuery__ScanToken * token = scan_result -> tokens [i ];
98+ if (token -> token == PG_QUERY__TOKEN__SQL_COMMENT || token -> token == PG_QUERY__TOKEN__C_COMMENT )
99+ result .comment_count ++ ;
100+ }
101+
102+ result .comments = malloc (result .comment_count * sizeof (PostgresDeparseComment * ));
103+ size_t comment_idx = 0 ;
104+ for (int i = 0 ; i < scan_result -> n_tokens ; i ++ )
105+ {
106+ PgQuery__ScanToken * token = scan_result -> tokens [i ];
107+ if (token -> token == PG_QUERY__TOKEN__SQL_COMMENT || token -> token == PG_QUERY__TOKEN__C_COMMENT )
108+ {
109+ size_t token_len = token -> end - token -> start ;
110+ PostgresDeparseComment * comment = malloc (sizeof (PostgresDeparseComment ));
111+
112+ /*
113+ * Set to the end of the prior non-comment token.
114+ *
115+ * We could alternatively set the match location to the start of the next
116+ * non-comment token instead (which would be more in line with emitting
117+ * that happens before the node once the location is matched).
118+ *
119+ * However, that doesn't work well in practice, specifically. RawStmt start
120+ * locations are 0, even when there are comments at the start of the statement.
121+ */
122+ comment -> match_location = prior_non_comment_end ;
123+
124+ comment -> newlines_before_comment = 0 ;
125+ comment -> newlines_after_comment = 0 ;
126+
127+ // Emit newlines before the comment if prior token was not a comment (otherwise that comment would have emitted them)
128+ if (!prior_token_was_comment )
129+ {
130+ for (int j = prior_token_end ; j < token -> start ; j ++ )
131+ {
132+ if (query [j ] == '\n' )
133+ comment -> newlines_before_comment ++ ;
134+ }
135+ }
136+
137+ if (i < scan_result -> n_tokens - 1 )
138+ {
139+ for (int j = token -> end ; j < scan_result -> tokens [i + 1 ]-> start ; j ++ )
140+ {
141+ if (query [j ] == '\n' )
142+ comment -> newlines_after_comment ++ ;
143+ }
144+ }
145+
146+ comment -> str = malloc (token_len + 1 );
147+ memcpy (comment -> str , & (query [token -> start ]), token_len );
148+ comment -> str [token_len ] = '\0' ;
149+
150+ result .comments [comment_idx ] = comment ;
151+ comment_idx ++ ;
152+ prior_token_was_comment = true;
153+ }
154+ else
155+ {
156+ prior_non_comment_end = token -> end ;
157+ prior_token_was_comment = false;
158+ }
159+ prior_token_end = token -> end ;
160+ }
161+
162+ pg_query__scan_result__free_unpacked (scan_result , NULL );
163+ pg_query_free_scan_result (scan_result_raw );
164+
165+ return result ;
166+ }
167+
168+ void pg_query_free_deparse_comments_result (PgQueryDeparseCommentsResult result )
169+ {
170+ for (int i = 0 ; i < result .comment_count ; i ++ )
171+ {
172+ free (result .comments [i ]-> str );
173+ free (result .comments [i ]);
174+ }
175+ if (result .comments != NULL )
176+ free (result .comments );
177+ }
0 commit comments