22
33PG_MODULE_MAGIC ;
44
5- static const char NEWLINE = '\n' ;
6- static const char DELIMITER = ',' ;
7- static const char DQUOTE = '"' ;
8- static const char CR = '\r' ;
5+ static const char NEWLINE = '\n' ;
6+ static const char DQUOTE = '"' ;
7+ static const char CR = '\r' ;
98
109typedef struct {
1110 StringInfoData accum_buf ;
@@ -14,17 +13,21 @@ typedef struct {
1413 TupleDesc tupdesc ;
1514} CsvAggState ;
1615
16+ static inline bool is_reserved (char c ) {
17+ return c == DQUOTE || c == NEWLINE || c == CR ;
18+ }
19+
1720// Any comma, quote, CR, LF requires quoting as per RFC https://www.ietf.org/rfc/rfc4180.txt
18- static inline bool needs_quote (const char * s , size_t n ) {
21+ static inline bool needs_quote (const char * s , size_t n , char delim ) {
1922 while (n -- ) {
2023 char c = * s ++ ;
21- if (c == DELIMITER || c == DQUOTE || c == NEWLINE || c == CR ) return true;
24+ if (c == delim || is_reserved ( c ) ) return true;
2225 }
2326 return false;
2427}
2528
26- static inline void csv_append_field (StringInfo buf , const char * s , size_t n ) {
27- if (!needs_quote (s , n )) {
29+ static inline void csv_append_field (StringInfo buf , const char * s , size_t n , char delim ) {
30+ if (!needs_quote (s , n , delim )) {
2831 appendBinaryStringInfo (buf , s , n );
2932 } else {
3033 appendStringInfoChar (buf , DQUOTE );
@@ -72,6 +75,10 @@ Datum csv_agg_transfn(PG_FUNCTION_ARGS) {
7275
7376 HeapTupleHeader next = PG_GETARG_HEAPTUPLEHEADER (1 );
7477
78+ char delim = PG_NARGS () >= 3 && !PG_ARGISNULL (2 ) ? PG_GETARG_CHAR (2 ) : ',' ;
79+
80+ if (is_reserved (delim )) ereport (ERROR , (errcode (ERRCODE_INVALID_PARAMETER_VALUE ), errmsg ("delimiter cannot be newline, carriage return or double quote" )));
81+
7582 // build header and cache tupdesc once
7683 if (!state -> header_done ) {
7784 TupleDesc tdesc = lookup_rowtype_tupdesc (HeapTupleHeaderGetTypeId (next ), HeapTupleHeaderGetTypMod (next ));
@@ -83,10 +90,10 @@ Datum csv_agg_transfn(PG_FUNCTION_ARGS) {
8390 continue ;
8491
8592 if (i > 0 ) // only append delimiter after the first value
86- appendStringInfoChar (& state -> accum_buf , DELIMITER );
93+ appendStringInfoChar (& state -> accum_buf , delim );
8794
8895 char * cstr = NameStr (att -> attname );
89- csv_append_field (& state -> accum_buf , cstr , strlen (cstr ));
96+ csv_append_field (& state -> accum_buf , cstr , strlen (cstr ), delim );
9097 }
9198
9299 appendStringInfoChar (& state -> accum_buf , NEWLINE );
@@ -119,12 +126,12 @@ Datum csv_agg_transfn(PG_FUNCTION_ARGS) {
119126 if (att -> attisdropped ) // pg always keeps dropped columns, guard against this
120127 continue ;
121128
122- if (i > 0 ) appendStringInfoChar (& state -> accum_buf , DELIMITER );
129+ if (i > 0 ) appendStringInfoChar (& state -> accum_buf , delim );
123130
124131 if (nulls [i ]) continue ; // empty field for NULL
125132
126133 char * cstr = datum_to_cstring (datums [i ], att -> atttypid );
127- csv_append_field (& state -> accum_buf , cstr , strlen (cstr ));
134+ csv_append_field (& state -> accum_buf , cstr , strlen (cstr ), delim );
128135 }
129136
130137 PG_RETURN_POINTER (state );
0 commit comments