copy.c 9.15 KB
Newer Older
Peter Eisentraut's avatar
Peter Eisentraut committed
1 2 3
/*
 * psql - the PostgreSQL interactive terminal
 *
4
 * Copyright 2000 by PostgreSQL Global Development Group
Peter Eisentraut's avatar
Peter Eisentraut committed
5
 *
6
 * $Header: /cvsroot/pgsql/src/bin/psql/copy.c,v 1.17 2001/02/10 02:31:28 tgl Exp $
Peter Eisentraut's avatar
Peter Eisentraut committed
7
 */
8
#include "postgres_fe.h"
9 10 11 12
#include "copy.h"

#include <errno.h>
#include <assert.h>
13
#include <signal.h>
14
#ifndef WIN32
Bruce Momjian's avatar
Bruce Momjian committed
15
#include <unistd.h>				/* for isatty */
16
#else
Bruce Momjian's avatar
Bruce Momjian committed
17
#include <io.h>					/* I think */
18 19
#endif

20
#include "libpq-fe.h"
21
#include "pqsignal.h"
22 23 24 25 26 27 28 29 30

#include "settings.h"
#include "common.h"
#include "stringutils.h"

#ifdef WIN32
#define strcasecmp(x,y) stricmp(x,y)
#endif

31
bool		copy_in_state;
32

33 34 35 36
/*
 * parse_slash_copy
 * -- parses \copy command line
 *
Peter Eisentraut's avatar
Peter Eisentraut committed
37
 * Accepted syntax: \copy [binary] table|"table" [with oids] from|to filename|'filename' [ using delimiters '<char>'] [ with null as 'string' ]
38 39 40 41 42
 * (binary is not here yet)
 *
 * returns a malloc'ed structure with the options, or NULL on parsing error
 */

Bruce Momjian's avatar
Bruce Momjian committed
43 44 45
struct copy_options
{
	char	   *table;
46
	char	   *file;			/* NULL = stdin/stdout */
Bruce Momjian's avatar
Bruce Momjian committed
47 48 49 50
	bool		from;
	bool		binary;
	bool		oids;
	char	   *delim;
51
	char	   *null;
52 53 54 55 56 57
};


static void
free_copy_options(struct copy_options * ptr)
{
Bruce Momjian's avatar
Bruce Momjian committed
58 59 60 61 62
	if (!ptr)
		return;
	free(ptr->table);
	free(ptr->file);
	free(ptr->delim);
63
	free(ptr->null);
Bruce Momjian's avatar
Bruce Momjian committed
64
	free(ptr);
65 66 67 68
}


static struct copy_options *
Peter Eisentraut's avatar
Peter Eisentraut committed
69
parse_slash_copy(const char *args)
70
{
Bruce Momjian's avatar
Bruce Momjian committed
71 72 73 74 75 76
	struct copy_options *result;
	char	   *line;
	char	   *token;
	bool		error = false;
	char		quote;

Peter Eisentraut's avatar
Peter Eisentraut committed
77 78 79 80 81 82 83
	if (args)
		line = xstrdup(args);
	else
	{
		psql_error("\\copy: arguments required\n");
		return NULL;
	}		
Bruce Momjian's avatar
Bruce Momjian committed
84 85 86

	if (!(result = calloc(1, sizeof(struct copy_options))))
	{
Peter Eisentraut's avatar
Peter Eisentraut committed
87
		psql_error("out of memory\n");
Bruce Momjian's avatar
Bruce Momjian committed
88 89 90
		exit(EXIT_FAILURE);
	}

91
	token = strtokx(line, " \t\n\r", "\"", '\\', &quote, NULL, pset.encoding);
Bruce Momjian's avatar
Bruce Momjian committed
92
	if (!token)
93
		error = true;
Bruce Momjian's avatar
Bruce Momjian committed
94 95
	else
	{
96
#ifdef NOT_USED
97 98
		/* this is not implemented yet */
		if (!quote && strcasecmp(token, "binary") == 0)
Bruce Momjian's avatar
Bruce Momjian committed
99 100
		{
			result->binary = true;
101
			token = strtokx(NULL, " \t\n\r", "\"", '\\', &quote, NULL, pset.encoding);
Bruce Momjian's avatar
Bruce Momjian committed
102 103 104 105
			if (!token)
				error = true;
		}
		if (token)
106
#endif
Bruce Momjian's avatar
Bruce Momjian committed
107
			result->table = xstrdup(token);
108 109 110
	}

#ifdef USE_ASSERT_CHECKING
Bruce Momjian's avatar
Bruce Momjian committed
111
	assert(error || result->table);
112 113
#endif

Bruce Momjian's avatar
Bruce Momjian committed
114 115
	if (!error)
	{
116
		token = strtokx(NULL, " \t\n\r", NULL, '\\', NULL, NULL, pset.encoding);
Bruce Momjian's avatar
Bruce Momjian committed
117
		if (!token)
118
			error = true;
Bruce Momjian's avatar
Bruce Momjian committed
119 120 121 122
		else
		{
			if (strcasecmp(token, "with") == 0)
			{
123
				token = strtokx(NULL, " \t\n\r", NULL, '\\', NULL, NULL, pset.encoding);
Bruce Momjian's avatar
Bruce Momjian committed
124 125 126 127 128 129 130
				if (!token || strcasecmp(token, "oids") != 0)
					error = true;
				else
					result->oids = true;

				if (!error)
				{
131
					token = strtokx(NULL, " \t\n\r", NULL, '\\', NULL, NULL, pset.encoding);
Bruce Momjian's avatar
Bruce Momjian committed
132 133 134 135 136 137 138 139 140 141 142
					if (!token)
						error = true;
				}
			}

			if (!error && strcasecmp(token, "from") == 0)
				result->from = true;
			else if (!error && strcasecmp(token, "to") == 0)
				result->from = false;
			else
				error = true;
143 144 145
		}
	}

Bruce Momjian's avatar
Bruce Momjian committed
146 147
	if (!error)
	{
148
		token = strtokx(NULL, " \t\n\r", "'", '\\', &quote, NULL, pset.encoding);
Bruce Momjian's avatar
Bruce Momjian committed
149 150
		if (!token)
			error = true;
151 152 153
		else if (!quote && (strcasecmp(token, "stdin") == 0 || strcasecmp(token, "stdout") == 0))
			result->file = NULL;
		else
Bruce Momjian's avatar
Bruce Momjian committed
154 155
			result->file = xstrdup(token);
	}
156

Bruce Momjian's avatar
Bruce Momjian committed
157 158
	if (!error)
	{
159
		token = strtokx(NULL, " \t\n\r", NULL, '\\', NULL, NULL, pset.encoding);
Bruce Momjian's avatar
Bruce Momjian committed
160 161
		if (token)
		{
162
			if (strcasecmp(token, "using") == 0)
Bruce Momjian's avatar
Bruce Momjian committed
163
			{
164
				token = strtokx(NULL, " \t\n\r", NULL, '\\', NULL, NULL, pset.encoding);
Bruce Momjian's avatar
Bruce Momjian committed
165 166 167 168
				if (!token || strcasecmp(token, "delimiters") != 0)
					error = true;
				else
				{
169
					token = strtokx(NULL, " \t\n\r", "'", '\\', NULL, NULL, pset.encoding);
Bruce Momjian's avatar
Bruce Momjian committed
170
					if (token)
171
					{
Bruce Momjian's avatar
Bruce Momjian committed
172
						result->delim = xstrdup(token);
173
						token = strtokx(NULL, " \t\n\r", NULL, '\\', NULL, NULL, pset.encoding);
174
					}
Bruce Momjian's avatar
Bruce Momjian committed
175 176 177 178
					else
						error = true;
				}
			}
179

180 181 182 183
			if (!error && token)
			{
				if (strcasecmp(token, "with") == 0)
				{
184
					token = strtokx(NULL, " \t\n\r", NULL, '\\', NULL, NULL, pset.encoding);
185 186 187 188
					if (!token || strcasecmp(token, "null") != 0)
						error = true;
					else
					{
189
						token = strtokx(NULL, " \t\n\r", NULL, '\\', NULL, NULL, pset.encoding);
190 191 192 193
						if (!token || strcasecmp(token, "as") != 0)
							error = true;
						else
						{
194
							token = strtokx(NULL, " \t\n\r", "'", '\\', NULL, NULL, pset.encoding);
195 196 197 198 199
							if (token)
								result->null = xstrdup(token);
						}
					}
				}
Peter Eisentraut's avatar
Peter Eisentraut committed
200 201
				else
					error = true;
202
			}
203 204 205
		}
	}

Bruce Momjian's avatar
Bruce Momjian committed
206
	free(line);
207

Bruce Momjian's avatar
Bruce Momjian committed
208 209
	if (error)
	{
210 211 212 213
		psql_error("\\copy: parse error at %s%s%s\n",
				   token ? "'" : "",
				   token ? token : "end of line",
				   token ? "'" : "");
Peter Eisentraut's avatar
Peter Eisentraut committed
214
		free_copy_options(result);
Bruce Momjian's avatar
Bruce Momjian committed
215 216
		return NULL;
	}
217
	else
Bruce Momjian's avatar
Bruce Momjian committed
218 219
		return result;
}
220 221 222 223 224 225 226



/*
 * Execute a \copy command (frontend copy). We have to open a file, then
 * submit a COPY query to the backend and either feed it data from the
 * file or route its response into the file.
Bruce Momjian's avatar
Bruce Momjian committed
227
 */
228
bool
Peter Eisentraut's avatar
Peter Eisentraut committed
229
do_copy(const char *args)
230
{
Bruce Momjian's avatar
Bruce Momjian committed
231 232 233 234 235
	char		query[128 + NAMEDATALEN];
	FILE	   *copystream;
	struct copy_options *options;
	PGresult   *result;
	bool		success;
236

Bruce Momjian's avatar
Bruce Momjian committed
237
	/* parse options */
Peter Eisentraut's avatar
Peter Eisentraut committed
238
	options = parse_slash_copy(args);
239

Bruce Momjian's avatar
Bruce Momjian committed
240 241
	if (!options)
		return false;
242

Bruce Momjian's avatar
Bruce Momjian committed
243 244
	strcpy(query, "COPY ");
	if (options->binary)
245
		strcat(query, "BINARY ");
246

Bruce Momjian's avatar
Bruce Momjian committed
247 248 249 250 251
	strcat(query, "\"");
	strncat(query, options->table, NAMEDATALEN);
	strcat(query, "\" ");
	if (options->oids)
		strcat(query, "WITH OIDS ");
252

Bruce Momjian's avatar
Bruce Momjian committed
253
	if (options->from)
254
		strcat(query, "FROM STDIN");
Bruce Momjian's avatar
Bruce Momjian committed
255
	else
256
		strcat(query, "TO STDOUT");
257 258


Bruce Momjian's avatar
Bruce Momjian committed
259 260
	if (options->delim)
	{
261
		strcat(query, " USING DELIMITERS '");
Bruce Momjian's avatar
Bruce Momjian committed
262 263 264
		strcat(query, options->delim);
		strcat(query, "'");
	}
265

266 267
	if (options->null)
	{
268 269 270 271
		strcat(query, " WITH NULL AS '");
		strcat(query, options->null);
		strcat(query, "'");
	}
272

Bruce Momjian's avatar
Bruce Momjian committed
273
	if (options->from)
274 275 276 277 278 279
	{
		if (options->file)
			copystream = fopen(options->file, "r");
		else
			copystream = stdin;
	}
Bruce Momjian's avatar
Bruce Momjian committed
280
	else
281 282 283 284 285 286
	{
		if (options->file)
			copystream = fopen(options->file, "w");
		else
			copystream = stdout;
	}
287

Bruce Momjian's avatar
Bruce Momjian committed
288 289
	if (!copystream)
	{
290 291
		psql_error("%s: %s\n",
				   options->file, strerror(errno));
Bruce Momjian's avatar
Bruce Momjian committed
292 293 294 295
		free_copy_options(options);
		return false;
	}

296
	result = PSQLexec(query);
Bruce Momjian's avatar
Bruce Momjian committed
297 298 299 300

	switch (PQresultStatus(result))
	{
		case PGRES_COPY_OUT:
301
			success = handleCopyOut(pset.db, copystream);
Bruce Momjian's avatar
Bruce Momjian committed
302 303
			break;
		case PGRES_COPY_IN:
304
			success = handleCopyIn(pset.db, copystream, NULL);
Bruce Momjian's avatar
Bruce Momjian committed
305 306 307 308 309
			break;
		case PGRES_NONFATAL_ERROR:
		case PGRES_FATAL_ERROR:
		case PGRES_BAD_RESPONSE:
			success = false;
310
			psql_error("\\copy: %s", PQerrorMessage(pset.db));
Bruce Momjian's avatar
Bruce Momjian committed
311 312 313
			break;
		default:
			success = false;
Peter Eisentraut's avatar
Peter Eisentraut committed
314
			psql_error("\\copy: unexpected response (%d)\n", PQresultStatus(result));
Bruce Momjian's avatar
Bruce Momjian committed
315 316 317
	}

	PQclear(result);
318

319 320
	if (copystream != stdout && copystream != stdin)
		fclose(copystream);
Bruce Momjian's avatar
Bruce Momjian committed
321 322
	free_copy_options(options);
	return success;
323 324 325 326 327 328 329
}


#define COPYBUFSIZ BLCKSZ


/*
330
 * handleCopyOut
331 332 333 334 335 336 337 338 339 340 341
 * receives data as a result of a COPY ... TO stdout command
 *
 * If you want to use COPY TO in your application, this is the code to steal :)
 *
 * conn should be a database connection that you just called COPY TO on
 * (and which gave you PGRES_COPY_OUT back);
 * copystream is the file stream you want the output to go to
 */
bool
handleCopyOut(PGconn *conn, FILE *copystream)
{
Bruce Momjian's avatar
Bruce Momjian committed
342 343 344
	bool		copydone = false;		/* haven't started yet */
	char		copybuf[COPYBUFSIZ];
	int			ret;
345

346
	assert(cancelConn);
347

Bruce Momjian's avatar
Bruce Momjian committed
348
	while (!copydone)
349
	{
Bruce Momjian's avatar
Bruce Momjian committed
350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372
		ret = PQgetline(conn, copybuf, COPYBUFSIZ);

		if (copybuf[0] == '\\' &&
			copybuf[1] == '.' &&
			copybuf[2] == '\0')
		{
			copydone = true;	/* we're at the end */
		}
		else
		{
			fputs(copybuf, copystream);
			switch (ret)
			{
				case EOF:
					copydone = true;
					/* FALLTHROUGH */
				case 0:
					fputc('\n', copystream);
					break;
				case 1:
					break;
			}
		}
373
	}
Bruce Momjian's avatar
Bruce Momjian committed
374
	fflush(copystream);
375
	ret = !PQendcopy(conn);
376 377
	cancelConn = NULL;
	return ret;
378 379 380 381 382
}



/*
383
 * handleCopyIn
384 385 386 387 388 389 390 391
 * receives data as a result of a COPY ... FROM stdin command
 *
 * Again, if you want to use COPY FROM in your application, copy this.
 *
 * conn should be a database connection that you just called COPY FROM on
 * (and which gave you PGRES_COPY_IN back);
 * copystream is the file stream you want the input to come from
 * prompt is something to display to request user input (only makes sense
Bruce Momjian's avatar
Bruce Momjian committed
392
 *	 if stdin is an interactive tty)
393 394 395
 */

bool
Bruce Momjian's avatar
Bruce Momjian committed
396
handleCopyIn(PGconn *conn, FILE *copystream, const char *prompt)
397
{
Bruce Momjian's avatar
Bruce Momjian committed
398 399 400 401 402
	bool		copydone = false;
	bool		firstload;
	bool		linedone;
	char		copybuf[COPYBUFSIZ];
	char	   *s;
403
	int			bufleft;
Bruce Momjian's avatar
Bruce Momjian committed
404
	int			c = 0;
405
	int			ret;
406
	unsigned int linecount=0;
Bruce Momjian's avatar
Bruce Momjian committed
407

408
#ifdef USE_ASSERT_CHECKING
409
	assert(copy_in_state);
410 411
#endif

412 413
	if (prompt)					/* disable prompt if not interactive */
	{
414
		if (!isatty(fileno(copystream)))
415 416 417
			prompt = NULL;
	}

Bruce Momjian's avatar
Bruce Momjian committed
418 419
	while (!copydone)
	{							/* for each input line ... */
420
		if (prompt)
421
		{
Bruce Momjian's avatar
Bruce Momjian committed
422 423 424 425 426
			fputs(prompt, stdout);
			fflush(stdout);
		}
		firstload = true;
		linedone = false;
427

Bruce Momjian's avatar
Bruce Momjian committed
428
		while (!linedone)
429
		{						/* for each bufferload in line ... */
Bruce Momjian's avatar
Bruce Momjian committed
430
			s = copybuf;
431
			for (bufleft = COPYBUFSIZ - 1; bufleft > 0; bufleft--)
Bruce Momjian's avatar
Bruce Momjian committed
432 433 434 435 436 437 438 439 440 441
			{
				c = getc(copystream);
				if (c == '\n' || c == EOF)
				{
					linedone = true;
					break;
				}
				*s++ = c;
			}
			*s = '\0';
442
			if (c == EOF && s == copybuf && firstload)
Bruce Momjian's avatar
Bruce Momjian committed
443 444 445
			{
				PQputline(conn, "\\.");
				copydone = true;
446 447
				if (pset.cur_cmd_interactive)
					puts("\\.");
Bruce Momjian's avatar
Bruce Momjian committed
448 449 450 451 452 453
				break;
			}
			PQputline(conn, copybuf);
			if (firstload)
			{
				if (!strcmp(copybuf, "\\."))
454
				{
Bruce Momjian's avatar
Bruce Momjian committed
455
					copydone = true;
456 457
					break;
				}
Bruce Momjian's avatar
Bruce Momjian committed
458 459
				firstload = false;
			}
460
		}
Bruce Momjian's avatar
Bruce Momjian committed
461
		PQputline(conn, "\n");
462
		linecount++;
463
	}
464
	ret = !PQendcopy(conn);
465
	copy_in_state = false;
466
	pset.lineno += linecount;
467
	return ret;
468
}