From f3f76040559902a9a429e029fd670a0180be9a34 Mon Sep 17 00:00:00 2001
From: Johanna Amann <johanna@icir.org>
Date: Wed, 12 Apr 2017 18:40:35 +0200
Subject: [PATCH] Add sql_addition config for writers, which allows UPSERT.

This commit adds a "sql_addition" configure option for the PostgreSQL
writer, which allows, e.g., to insert an UPSERT option after the insert
call that is constructed by the plugin.

For a complete example, see the write-upsert testcase, which uses this
command to create a filter:

local filter: Log::Filter = [$name="postgres", $path="testtable", $writer=Log::WRITER_POSTGRESQL, $config=table(["dbname"]="testdb", ["port"]="7772", ["sql_addition"]="ON CONFLICT (i) DO UPDATE SET s=EXCLUDED.s")];
---
 src/PostgresWriter.cc                        |  8 ++-
 src/PostgresWriter.h                         |  2 +-
 tests/Baseline/postgres.write-upsert/ssh.out |  3 +
 tests/postgres/write-upsert.bro              | 59 ++++++++++++++++++++
 4 files changed, 68 insertions(+), 4 deletions(-)
 create mode 100644 tests/Baseline/postgres.write-upsert/ssh.out
 create mode 100644 tests/postgres/write-upsert.bro

diff --git a/src/PostgresWriter.cc b/src/PostgresWriter.cc
index 3508826..49defc2 100644
--- a/src/PostgresWriter.cc
+++ b/src/PostgresWriter.cc
@@ -99,7 +99,7 @@ string PostgreSQL::GetTableType(int arg_type, int arg_subtype)
 }
 
 // preformat the insert string that we only need to create once during our lifetime
-bool PostgreSQL::CreateInsert(int num_fields, const Field* const * fields)
+bool PostgreSQL::CreateInsert(int num_fields, const Field* const * fields, std::string add_string)
 	{
 	string names = "INSERT INTO "+table+" ( ";
 	string values("VALUES (");
@@ -120,7 +120,7 @@ bool PostgreSQL::CreateInsert(int num_fields, const Field* const * fields)
 		values += "$" + std::to_string(i+1);
 		}
 
-	insert = names + ") " + values + ");";
+	insert = names + ") " + values + ") " + add_string + ";";
 
 	return true;
 	}
@@ -186,6 +186,8 @@ bool PostgreSQL::DoInit(const WriterInfo& info, int num_fields,
 			conninfo += Fmt(" port = %d", default_port);
 		}
 
+	string add_string = LookupParam(info, "sql_addition");
+
 	string errorhandling = LookupParam(info, "continue_on_errors");
 	if ( !errorhandling.empty() && errorhandling == "T" )
 		ignore_errors = true;
@@ -238,7 +240,7 @@ bool PostgreSQL::DoInit(const WriterInfo& info, int num_fields,
 		return false;
 		}
 
-	return CreateInsert(num_fields, fields);
+	return CreateInsert(num_fields, fields, add_string);
 	}
 
 bool PostgreSQL::DoFlush(double network_time)
diff --git a/src/PostgresWriter.h b/src/PostgresWriter.h
index c724cc9..c9807b4 100644
--- a/src/PostgresWriter.h
+++ b/src/PostgresWriter.h
@@ -41,7 +41,7 @@ class PostgreSQL : public WriterBackend {
 	string EscapeIdentifier(const char* identifier);
 	std::tuple<bool, string, int> CreateParams(const threading::Value* val);
 	string GetTableType(int, int);
-	bool CreateInsert(int num_fields, const threading::Field* const* fields);
+	bool CreateInsert(int num_fields, const threading::Field* const* fields, const std::string add_string = "");
 
 	PGconn *conn;
 
diff --git a/tests/Baseline/postgres.write-upsert/ssh.out b/tests/Baseline/postgres.write-upsert/ssh.out
new file mode 100644
index 0000000..798df00
--- /dev/null
+++ b/tests/Baseline/postgres.write-upsert/ssh.out
@@ -0,0 +1,3 @@
+i|s
+-42|hurz2
+(1 row)
diff --git a/tests/postgres/write-upsert.bro b/tests/postgres/write-upsert.bro
new file mode 100644
index 0000000..efe5517
--- /dev/null
+++ b/tests/postgres/write-upsert.bro
@@ -0,0 +1,59 @@
+# @TEST-SERIALIZE: postgres
+# @TEST-EXEC: initdb postgres
+# @TEST-EXEC: perl -pi.bak -E "s/#port =.*/port = 7772/;" postgres/postgresql.conf
+# @TEST-EXEC: pg_ctl start -D postgres -l serverlog
+# @TEST-EXEC: sleep 5
+# @TEST-EXEC: createdb -p 7772 testdb
+# @TEST-EXEC: psql -p 7772 testdb < create.sql
+# @TEST-EXEC: bro %INPUT || true
+# @TEST-EXEC: echo "select * from testtable" | psql -A -p 7772 testdb >ssh.out 2>&1 || true
+# @TEST-EXEC: pg_ctl stop -D postgres -m fast
+# @TEST-EXEC: btest-diff ssh.out
+
+@TEST-START-FILE create.sql
+create table testtable (
+i integer not null unique,
+s varchar not null unique);
+@TEST-END-FILE
+
+# Test all possible types.
+
+module SSHTest;
+
+export {
+	redef enum Log::ID += { LOG };
+
+	type Log: record {
+		i: int;
+		s: string;
+	} &log;
+}
+
+function foo(i : count) : string
+	{
+	if ( i > 0 )
+		return "Foo";
+	else
+		return "Bar";
+	}
+
+event bro_init()
+{
+	Log::create_stream(SSHTest::LOG, [$columns=Log]);
+	local filter: Log::Filter = [$name="postgres", $path="testtable", $writer=Log::WRITER_POSTGRESQL, $config=table(["dbname"]="testdb", ["port"]="7772", ["sql_addition"]="ON CONFLICT (i) DO UPDATE SET s=EXCLUDED.s")];
+	Log::add_filter(SSHTest::LOG, filter);
+
+	local empty_set: set[string];
+	local empty_vector: vector of string;
+
+	Log::write(SSHTest::LOG, [
+		$i=-42,
+		$s="hurz"
+		]);
+
+	Log::write(SSHTest::LOG, [
+		$i=-42,
+		$s="hurz2"
+		]);
+}
+