File: qbuilder.c

package info (click to toggle)
pgq 3.2.6-8
  • links: PTS, VCS
  • area: main
  • in suites: bullseye, buster, sid
  • size: 836 kB
  • sloc: sql: 3,412; ansic: 1,990; python: 302; makefile: 89; sh: 2
file content (147 lines) | stat: -rw-r--r-- 3,259 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147

#include <postgres.h>
#include <executor/spi.h>

#include "qbuilder.h"
#include "parsesql.h"

/* import standard_conforming_strings */
#if PG_VERSION_NUM >= 80500
#include <parser/parser.h>
#else
#ifndef PGDLLIMPORT
#define PGDLLIMPORT DLLIMPORT
#endif
extern PGDLLIMPORT bool standard_conforming_strings;
#endif

/* create QB in right context */
struct QueryBuilder *qb_create(const struct QueryBuilderOps *ops, MemoryContext ctx)
{
	struct QueryBuilder *q;

	q = MemoryContextAllocZero(ctx, sizeof(*q));
	q->op = ops;
	q->stdstr = standard_conforming_strings;

	q->maxargs = 8;
	q->arg_map = MemoryContextAlloc(ctx, q->maxargs * sizeof(int));

	/* default size too large? */
	q->sql.maxlen = 64;
	q->sql.data = MemoryContextAlloc(ctx,  q->sql.maxlen);
	q->sql.data[0] = 0;
	return q;
}

/* add fragment without parsing */
void qb_add_raw(struct QueryBuilder *q, const char *str, int len)
{
	if (len < 0)
		len = strlen(str);
	appendBinaryStringInfo(&q->sql, str, len);
}

/* the ident may or may not be argument reference */
static void qb_handle_ident(struct QueryBuilder *q, const char *ident, int len, void *arg)
{
	int real_idx;
	int local_idx = -1, i;
	char abuf[32];

	/* is argument reference? */
	real_idx = q->op->name_lookup(arg, ident, len);
	if (real_idx < 0) {
		qb_add_raw(q, ident, len);
		return;
	}

	/* already referenced? */
	for (i = 0; i < q->nargs; i++) {
		if (q->arg_map[i] == real_idx) {
			local_idx = i;
			break;
		}
	}

	/* new reference? */
	if (local_idx < 0) {
		if (q->nargs >= FUNC_MAX_ARGS)
			elog(ERROR, "Too many args");
		if (q->nargs >= q->maxargs) {
			q->arg_map = repalloc(q->arg_map, q->maxargs * 2 * sizeof(int));
			q->maxargs *= 2;
		}
		local_idx = q->nargs++;
		q->arg_map[local_idx] = real_idx;
	}

	/* add $n to query */
	snprintf(abuf, sizeof(abuf), "$%d", local_idx + 1);
	return qb_add_raw(q, abuf, strlen(abuf));
}

/* add fragment with parsing - argument references are replaced with $n */
void qb_add_parse(struct QueryBuilder *q, const char *sql, void *arg)
{
	int tlen, tok;

	/* tokenize sql, pick out argument references */
	while (1) {
		tok = sql_tokenizer(sql, &tlen, q->stdstr);
		if (!tok)
			break;
		if (tok < 0)
			elog(ERROR, "QB: syntax error");
		if (tok == T_WORD) {
			qb_handle_ident(q, sql, tlen, arg);
		} else {
			qb_add_raw(q, sql, tlen);
		}
		sql += tlen;
	}
}

/* prepare */
void qb_prepare(struct QueryBuilder *q, void *arg)
{
	Oid types[FUNC_MAX_ARGS];
	void *plan;
	int i;

	for (i = 0; i < q->nargs; i++)
		types[i] = q->op->type_lookup(arg, q->arg_map[i]);

	plan = SPI_prepare(q->sql.data, q->nargs, types);
	q->plan = SPI_saveplan(plan);
}

/* lookup values and run plan.  returns result from SPI_execute_plan()  */
int qb_execute(struct QueryBuilder *q, void *arg)
{
	Datum values[FUNC_MAX_ARGS];
	char nulls[FUNC_MAX_ARGS];
	int i;

	if (!q->plan)
		elog(ERROR, "QB: query not prepared yet");

	for (i = 0; i < q->nargs; i++) {
		bool isnull = false;
		values[i] = q->op->value_lookup(arg, q->arg_map[i], &isnull);
		nulls[i] = isnull ? 'n' : ' ';
	}
	return SPI_execute_plan(q->plan, values, nulls, true, 0);
}

void qb_free(struct QueryBuilder *q)
{
	if (!q)
		return;
	if (q->plan)
		SPI_freeplan(q->plan);
	if (q->sql.data)
		pfree(q->sql.data);
	pfree(q);
}