Использование кастомных форматов и протоколов для внешних таблиц

#include "postgres.h"
#include "fmgr.h"
#include "funcapi.h"

#include "access/extprotocol.h"
#include "catalog/pg_proc.h"
#include "commands/defrem.h"
#include "utils/array.h"
#include "utils/builtins.h"
#include "utils/memutils.h"

#include "catalog/pg_exttable.h"


typedef struct DemoUri
{
	char	   *protocol;
	char	   *path;

}	DemoUri;

static DemoUri *ParseDemoUri(const char *uri_str);
static void FreeDemoUri(DemoUri* uri);


/* Do the module magic dance */
PG_MODULE_MAGIC;
PG_FUNCTION_INFO_V1(demoprot_export);
PG_FUNCTION_INFO_V1(demoprot_import);
PG_FUNCTION_INFO_V1(demoprot_validate_urls);

Datum demoprot_export(PG_FUNCTION_ARGS);
Datum demoprot_import(PG_FUNCTION_ARGS);
Datum demoprot_validate_urls(PG_FUNCTION_ARGS);


typedef struct {
	char	  *url;
	char	  *filename;
	FILE	  *file;
} extprotocol_t;

static void check_ext_options(const FunctionCallInfo fcinfo)
{
	ListCell *cell;
	Relation rel = EXTPROTOCOL_GET_RELATION(fcinfo);
	ExtTableEntry *exttbl = GetExtTableEntry(rel->rd_id);
	List *options = exttbl->options;

	foreach(cell, options) {
		DefElem *def = (DefElem *) lfirst(cell);
		char *key = def->defname;
		char *value = defGetString(def);

		if (key && strcasestr(key, "database") && !strcasestr(value, "greengage")) {
				ereport(ERROR, (0, errmsg("This is greengage.")));
		}
	}
}

/*
 * Import data into GPDB.
 */
Datum
demoprot_import(PG_FUNCTION_ARGS)
{
	extprotocol_t   *myData;
	char			*data;
	int				 datlen;
	size_t			 nread = 0;

	/* Must be called via the external table format manager */
	if (!CALLED_AS_EXTPROTOCOL(fcinfo))
		elog(ERROR, "extprotocol_import: not called by external protocol manager");

	/* Get our internal description of the protocol */
	myData = (extprotocol_t *) EXTPROTOCOL_GET_USER_CTX(fcinfo);

	if(EXTPROTOCOL_IS_LAST_CALL(fcinfo))
	{
		/* we're done receiving data. close our connection */
		if(myData && myData->file)
			if(fclose(myData->file))
				ereport(ERROR,
						(errcode_for_file_access(),
						 errmsg("could not close file \"%s\": %m",
								 myData->filename)));

		PG_RETURN_INT32(0);
	}

	if (myData == NULL)
	{
		/* first call. do any desired init */

		const char	*p_name = "demoprot";
		DemoUri		*parsed_url;
		char		*url = EXTPROTOCOL_GET_URL(fcinfo);

		myData 			 = palloc(sizeof(extprotocol_t));

		myData->url 	 = pstrdup(url);
		parsed_url 		 = ParseDemoUri(myData->url);
		myData->filename = pstrdup(parsed_url->path);

		if(strcasecmp(parsed_url->protocol, p_name) != 0)
			elog(ERROR, "internal error: demoprot called with a different protocol (%s)",
						parsed_url->protocol);

		/* An example of checking options */
		check_ext_options(fcinfo);

		FreeDemoUri(parsed_url);

		/* open the destination file (or connect to remote server in other cases) */
		myData->file = fopen(myData->filename, "r");

		if (myData->file == NULL)
			ereport(ERROR,
					(errcode_for_file_access(),
					 errmsg("demoprot_import: could not open file \"%s\" for reading: %m",
							myData->filename)));

		EXTPROTOCOL_SET_USER_CTX(fcinfo, myData);
	}

	/* =======================================================================
	 *                            DO THE IMPORT
	 * ======================================================================= */

	data 	= EXTPROTOCOL_GET_DATABUF(fcinfo);
	datlen 	= EXTPROTOCOL_GET_DATALEN(fcinfo);

	if(datlen > 0)
	{
		nread = fread(data, 1, datlen, myData->file);
		if (ferror(myData->file))
			ereport(ERROR,
					(errcode_for_file_access(),
					 errmsg("demoprot_import: could not write to file \"%s\": %m",
							 myData->filename)));
	}


	PG_RETURN_INT32((int)nread);
}

/*
 * Export data out of GPDB.
 */
Datum
demoprot_export(PG_FUNCTION_ARGS)
{
	extprotocol_t   *myData;
	char			*data;
	int				 datlen;
	size_t			 wrote = 0;

	/* Must be called via the external table format manager */
	if (!CALLED_AS_EXTPROTOCOL(fcinfo))
		elog(ERROR, "extprotocol_export: not called by external protocol manager");

	/* Get our internal description of the protocol */
	myData = (extprotocol_t *) EXTPROTOCOL_GET_USER_CTX(fcinfo);

	if(EXTPROTOCOL_IS_LAST_CALL(fcinfo))
	{
		/* we're done sending data. close our connection */
		if(myData && myData->file)
			if(fclose(myData->file))
				ereport(ERROR,
						(errcode_for_file_access(),
						 errmsg("could not close file \"%s\": %m",
								 myData->filename)));

		PG_RETURN_INT32(0);
	}

	if (myData == NULL)
	{
		/* first call. do any desired init */

		const char	*p_name = "demoprot";
		DemoUri		*parsed_url;
		char		*url = EXTPROTOCOL_GET_URL(fcinfo);

		myData 			 = palloc(sizeof(extprotocol_t));

		myData->url 	 = pstrdup(url);
		parsed_url 		 = ParseDemoUri(myData->url);
		myData->filename = pstrdup(parsed_url->path);

		if(strcasecmp(parsed_url->protocol, p_name) != 0)
			elog(ERROR, "internal error: demoprot called with a different protocol (%s)",
						parsed_url->protocol);

		FreeDemoUri(parsed_url);

		/* open the destination file (or connect to remote server in other cases) */
		myData->file = fopen(myData->filename, "a");

		if (myData->file == NULL)
			ereport(ERROR,
					(errcode_for_file_access(),
					 errmsg("demoprot_export: could not open file \"%s\" for writing: %m",
							myData->filename)));

		EXTPROTOCOL_SET_USER_CTX(fcinfo, myData);
	}

	/* =======================================================================
	 *                            DO THE EXPORT
	 * ======================================================================= */

	data 	= EXTPROTOCOL_GET_DATABUF(fcinfo);
	datlen 	= EXTPROTOCOL_GET_DATALEN(fcinfo);

	if(datlen > 0)
	{
		wrote = fwrite(data, 1, datlen, myData->file);
		if (ferror(myData->file))
			ereport(ERROR,
					(errcode_for_file_access(),
					 errmsg("demoprot_import: could not read from file \"%s\": %m",
							 myData->filename)));
	}

	PG_RETURN_INT32((int)wrote);
}

Datum
demoprot_validate_urls(PG_FUNCTION_ARGS)
{
	int					nurls;
	int					i;
	ValidatorDirection	direction;

	/* Must be called via the external table format manager */
	if (!CALLED_AS_EXTPROTOCOL_VALIDATOR(fcinfo))
		elog(ERROR, "demoprot_validate_urls: not called by external protocol manager");

	nurls 		= EXTPROTOCOL_VALIDATOR_GET_NUM_URLS(fcinfo);
	direction 	= EXTPROTOCOL_VALIDATOR_GET_DIRECTION(fcinfo);

	/*
	 * Dumb example 1: search each url for a substring
	 * we don't want to be used in a url. in this example
	 * it's 'secured_directory'.
	 */
	for (i = 1 ; i <= nurls ; i++)
	{
		char *url = EXTPROTOCOL_VALIDATOR_GET_NTH_URL(fcinfo, i);

		if (strstr(url, "secured_directory") != 0)
		{
            ereport(ERROR,
                    (errcode(ERRCODE_PROTOCOL_VIOLATION),
                     errmsg("using 'secured_directory' in a url isn't allowed ")));
		}
	}

	/*
	 * Dumb example 2: set a limit on the number of urls
	 * used. In this example we limit readable external
	 * tables that use our protocol to 2 urls max.
	 */
	if(direction == EXT_VALIDATE_READ && nurls > 2)
	{
        ereport(ERROR,
                (errcode(ERRCODE_PROTOCOL_VIOLATION),
                 errmsg("more than 2 urls aren't allowed in this protocol ")));
	}

	PG_RETURN_VOID();
}

/* --- utility functions --- */

static
DemoUri *ParseDemoUri(const char *uri_str)
{
	DemoUri	   *uri = (DemoUri *) palloc0(sizeof(DemoUri));
	int			protocol_len;

 	uri->path = NULL;
 	uri->protocol = NULL;

	/*
	 * parse protocol
	 */
	char *post_protocol = strstr(uri_str, "://");

	if(!post_protocol)
	{
		ereport(ERROR,
				(errcode(ERRCODE_SYNTAX_ERROR),
				 errmsg("invalid demo prot URI \'%s\'", uri_str)));
	}

	protocol_len = post_protocol - uri_str;
	uri->protocol = (char *) palloc0 (protocol_len + 1);
	strncpy(uri->protocol, uri_str, protocol_len);

	/* make sure there is more to the uri string */
	if (strlen(uri_str) <= protocol_len)
		ereport(ERROR,
				(errcode(ERRCODE_SYNTAX_ERROR),
		errmsg("invalid demo prot URI \'%s\' : missing path", uri_str)));

	/*
	 * parse path
	 */
	uri->path = pstrdup(uri_str + protocol_len + strlen("://"));

	return uri;
}

static
void FreeDemoUri(DemoUri *uri)
{
	if (uri->path)
		pfree(uri->path);
	if (uri->protocol)
		pfree(uri->protocol);

	pfree(uri);
}

#include "postgres.h"
#include "fmgr.h"
#include "funcapi.h"

#include "access/formatter.h"
#include "catalog/pg_proc.h"
#include "utils/builtins.h"
#include "utils/memutils.h"
#include "utils/typcache.h"
#include "utils/syscache.h"

/* Do the module magic dance */
PG_MODULE_MAGIC;
PG_FUNCTION_INFO_V1(formatter_export);
PG_FUNCTION_INFO_V1(formatter_import);

Datum formatter_export(PG_FUNCTION_ARGS);
Datum formatter_import(PG_FUNCTION_ARGS);


typedef struct {
	int        ncols;
	Datum     *values;
	bool      *nulls;
	int        buflen;
	bytea     *buffer;
} format_t;


/*
 * Maximum size string to support, affects allocation size of the tuple buffer.
 * Only used for variable length strings.  For strings with a declared typmod
 * we allow that size even if it is larger than this.
 */
#define MAX_FORMAT_STRING 4096

/*
 * Our format converts all NULLs to real values, for floats that value is NaN
 */
#define NULL_FLOAT8_VALUE get_float8_nan()


Datum
formatter_export(PG_FUNCTION_ARGS)
{
	HeapTupleHeader		rec	= PG_GETARG_HEAPTUPLEHEADER(0);
	TupleDesc           tupdesc;
	HeapTupleData		tuple;
	int                 ncolumns = 0;
	format_t           *myData;
	char               *data;
	int                 datlen;
	int                 i;

	/* Must be called via the external table format manager */
	if (!CALLED_AS_FORMATTER(fcinfo))
		elog(ERROR, "formatter_export: not called by format manager");

	tupdesc = FORMATTER_GET_TUPDESC(fcinfo);

	/* Get our internal description of the formatter */
	ncolumns = tupdesc->natts;
	myData = (format_t *) FORMATTER_GET_USER_CTX(fcinfo);
	if (myData == NULL)
	{
		myData          = palloc(sizeof(format_t));
		myData->ncols   = ncolumns;
		myData->values  = palloc(sizeof(Datum) * ncolumns);
		myData->nulls   = palloc(sizeof(bool) * ncolumns);

		/* Determine required buffer size */
		myData->buflen = 0;
		for (i = 0; i < ncolumns; i++)
		{
			Oid   type   = tupdesc->attrs[i]->atttypid;
			int32 typmod = tupdesc->attrs[i]->atttypmod;

			/* Don't know how to format dropped columns, error for now */
			if (tupdesc->attrs[i]->attisdropped)
				elog(ERROR, "formatter_export: dropped columns");

			switch (type)
			{
				case FLOAT8OID:
				{
					myData->buflen += sizeof(double);
					break;
				}

				case VARCHAROID:
				case BPCHAROID:
				case TEXTOID:
				{
					myData->buflen += (typmod > 0) ? typmod : MAX_FORMAT_STRING;
					break;
				}

				default:
				{
					elog(ERROR, "formatter_export error: unsupported data type");
					break;
				}
			}
		}

		myData->buflen = Max(128, myData->buflen);  /* allocate at least 128 bytes */
		myData->buffer = palloc(myData->buflen + VARHDRSZ);

		FORMATTER_SET_USER_CTX(fcinfo, myData);
	}
	if (myData->ncols != ncolumns)
		elog(ERROR, "formatter_export: unexpected change of output record type");

	/* break the input tuple into fields */
	tuple.t_len = HeapTupleHeaderGetDatumLength(rec);
	ItemPointerSetInvalid(&(tuple.t_self));
	tuple.t_data = rec;
	heap_deform_tuple(&tuple, tupdesc, myData->values, myData->nulls);

	datlen = 0;
	data = VARDATA(myData->buffer);


	/* =======================================================================
	 *                            MAIN FORMATTING CODE
	 *
	 * Currently this code assumes:
	 *  - Homogoneos hardware => No need to convert data to network byte order
	 *  - Support for TEXT/VARCHAR/BPCHAR/FLOAT8 only
	 *  - Length Prefixed strings
	 *  - No end of record tags, checksums, or optimizations for alignment.
	 *  - NULL values are cast to some sensible default value (NaN, "")
	 *
	 * ======================================================================= */
	for (i = 0; i < ncolumns; i++)
	{
		Oid	  type    = tupdesc->attrs[i]->atttypid;
		int   typmod  = tupdesc->attrs[i]->atttypmod;
		Datum val     = myData->values[i];
		bool  nul     = myData->nulls[i];

		switch (type)
		{
			case FLOAT8OID:
			{
				float8 value;

				if (datlen + sizeof(value) >= myData->buflen)
					elog(ERROR, "formatter_export: buffer too small");

				if (nul)
					value = NULL_FLOAT8_VALUE;
				else
					value = DatumGetFloat8(val);

				memcpy(&data[datlen], &value, sizeof(value));
				datlen += sizeof(value);
				break;
			}

			case TEXTOID:
			case VARCHAROID:
			case BPCHAROID:
			{
				text  *str;
				int32  len;

				if (nul)
				{
					str = NULL;
					len = 0;
				}
				else
				{
					str = DatumGetTextP(val);
					len = VARSIZE(str) - VARHDRSZ;
					if (typmod < 0)
						len  = Min(len, MAX_FORMAT_STRING);
				}

				if (datlen + sizeof(len) + len >= myData->buflen)
					elog(ERROR, "formatter_export: buffer too small");
				memcpy(&data[datlen], &len, sizeof(len));
				datlen += sizeof(len);

				if (len > 0)
				{
					memcpy(&data[datlen], VARDATA(str), len);
					datlen += len;
				}
				break;
			}

			default:
				elog(ERROR, "formatter_export: unsupported datatype");
				break;
		}
	}
	/* ======================================================================= */

	SET_VARSIZE(myData->buffer, datlen + VARHDRSZ);
	PG_RETURN_BYTEA_P(myData->buffer);
}

Datum
formatter_import(PG_FUNCTION_ARGS)
{
	HeapTuple			tuple;
	TupleDesc           tupdesc;
	MemoryContext 		m, oldcontext;
	format_t           *myData;
	char               *data_buf;
	int                 ncolumns = 0;
	int			  		data_cur;
	int                 data_len;
	int                 i;

	/* Must be called via the external table format manager */
	if (!CALLED_AS_FORMATTER(fcinfo))
		elog(ERROR, "formatter_import: not called by format manager");

	tupdesc = FORMATTER_GET_TUPDESC(fcinfo);

	/* Get our internal description of the formatter */
	ncolumns = tupdesc->natts;
	myData = (format_t *) FORMATTER_GET_USER_CTX(fcinfo);

	if (myData == NULL)
	{

		myData          = palloc(sizeof(format_t));
		myData->ncols   = ncolumns;
		myData->values  = palloc(sizeof(Datum) * ncolumns);
		myData->nulls   = palloc(sizeof(bool) * ncolumns);

		/* misc verification */
		for (i = 0; i < ncolumns; i++)
		{
			Oid   type   = tupdesc->attrs[i]->atttypid;
			//int32 typmod = tupdesc->attrs[i]->atttypmod;

			/* Don't know how to format dropped columns, error for now */
			if (tupdesc->attrs[i]->attisdropped)
				elog(ERROR, "formatter_import: dropped columns");

			switch (type)
			{
				case FLOAT8OID:
				case VARCHAROID:
				case BPCHAROID:
				case TEXTOID:
					break;

				default:
				{
					elog(ERROR, "formatter_import error: unsupported data type");
					break;
				}
			}
		}

		FORMATTER_SET_USER_CTX(fcinfo, myData);

	}
	if (myData->ncols != ncolumns)
		elog(ERROR, "formatter_import: unexpected change of output record type");

	/* get our input data buf and number of valid bytes in it */
	data_buf = FORMATTER_GET_DATABUF(fcinfo);
	data_len = FORMATTER_GET_DATALEN(fcinfo);
	data_cur = FORMATTER_GET_DATACURSOR(fcinfo);

	/* start clean */
	MemSet(myData->values, 0, ncolumns * sizeof(Datum));
	MemSet(myData->nulls, true, ncolumns * sizeof(bool));

	/* =======================================================================
	 *                            MAIN FORMATTING CODE
	 *
	 * Currently this code assumes:
	 *  - Homogoneos hardware => No need to convert data to network byte order
	 *  - Support for TEXT/VARCHAR/BPCHAR/FLOAT8 only
	 *  - Length Prefixed strings
	 *  - No end of record tags, checksums, or optimizations for alignment.
	 *  - NULL values are cast to some sensible default value (NaN, "")
	 *
	 * ======================================================================= */
	m = FORMATTER_GET_PER_ROW_MEM_CTX(fcinfo);
	oldcontext = MemoryContextSwitchTo(m);

	for (i = 0; i < ncolumns; i++)
	{
		Oid		type    	= tupdesc->attrs[i]->atttypid;
		//int	typmod		= tupdesc->attrs[i]->atttypmod;
		int		remaining	= 0;
		int		attr_len 	= 0;

		remaining = data_len - data_cur;

		switch (type)
		{
			case FLOAT8OID:
			{
				float8 value;

				attr_len = sizeof(value);

				if (remaining < attr_len)
				{
					MemoryContextSwitchTo(oldcontext);
					FORMATTER_RETURN_NOTIFICATION(fcinfo, FMT_NEED_MORE_DATA);
				}

				memcpy(&value, data_buf + data_cur, attr_len);

				if(value != NULL_FLOAT8_VALUE)
				{
					myData->nulls[i] = false;
					myData->values[i] = Float8GetDatum(value);
				}

				/* TODO: check for nan? */

				break;
			}

			case TEXTOID:
			case VARCHAROID:
			case BPCHAROID:
			{
				text*	value;
				int32	len;
				bool	nextlen = remaining >= sizeof(len);

				if (nextlen)
				{
					memcpy(&len, data_buf + data_cur, sizeof(len));

					if (len < 0)
						elog(ERROR, "invalid length of varlen datatype: %d",
									len);
				}

				/* if len or data bytes don't exist in this buffer, return */
				if (!nextlen || (nextlen && (remaining - sizeof(len) < len)))
				{
					MemoryContextSwitchTo(oldcontext);
					FORMATTER_RETURN_NOTIFICATION(fcinfo, FMT_NEED_MORE_DATA);
				}

				if (len > 0)
				{
					value = (text *) palloc(len + VARHDRSZ);
					SET_VARSIZE(value, len + VARHDRSZ);

					memcpy(VARDATA(value), data_buf + data_cur + sizeof(len), len);

					myData->nulls[i] = false;
					myData->values[i] = PointerGetDatum(value);
				}

				attr_len = len + sizeof(len);

				break;
			}

			default:
				elog(ERROR, "formatter_import: unsupported datatype");
				break;
		}

		/* add byte length of last attribute to the temporary cursor */
		data_cur += attr_len;

	}
	/* ======================================================================= */

	MemoryContextSwitchTo(oldcontext);

	FORMATTER_SET_DATACURSOR(fcinfo, data_cur);

	tuple = heap_form_tuple(tupdesc, myData->values, myData->nulls);

	/* hack... pass tuple here. don't free prev tuple - the executor does it  */
	((FormatterData*) fcinfo->context)->fmt_tuple = tuple;

	FORMATTER_RETURN_TUPLE(tuple);
}