Skip to content

Commit

Permalink
Add sql_addition config for writers, which allows UPSERT.
Browse files Browse the repository at this point in the history
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")];
  • Loading branch information
0xxon committed Apr 12, 2017
1 parent f5a8ca2 commit f3f7604
Show file tree
Hide file tree
Showing 4 changed files with 68 additions and 4 deletions.
8 changes: 5 additions & 3 deletions src/PostgresWriter.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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 (");
Expand All @@ -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;
}
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion src/PostgresWriter.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
3 changes: 3 additions & 0 deletions tests/Baseline/postgres.write-upsert/ssh.out
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
i|s
-42|hurz2
(1 row)
59 changes: 59 additions & 0 deletions tests/postgres/write-upsert.bro
Original file line number Diff line number Diff line change
@@ -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"
]);
}

0 comments on commit f3f7604

Please sign in to comment.