ctdb-tools/ctdb: New ptrans command

Also add test.

Signed-off-by: Martin Schwenke <martin@meltin.net>
Pair-programmed-with: Amitay Isaacs <amitay@gmail.com>
Reviewed-by: Michael Adam <obnox@samba.org>

3 files changed, 278 insertions, 0 deletions

diff --git a/ctdb/include/ctdb_client.h b/ctdb/include/ctdb_client.h
index 53b082978ca..d3084273ea5 100644
--- a/ctdb/include/ctdb_client.h
+++ b/ctdb/include/ctdb_client.h
@@ -550,6 +550,7 @@ int ctdb_transaction_fetch(struct ctdb_transaction_handle *h,
 int ctdb_transaction_store(struct ctdb_transaction_handle *h,
 			   TDB_DATA key, TDB_DATA data);
 int ctdb_transaction_commit(struct ctdb_transaction_handle *h);
+int ctdb_transaction_cancel(struct ctdb_transaction_handle *h);
 
 int ctdb_ctrl_recd_ping(struct ctdb_context *ctdb);
 
diff --git a/ctdb/tests/simple/55_ctdb_ptrans.sh b/ctdb/tests/simple/55_ctdb_ptrans.sh
new file mode 100755
index 00000000000..14de67efa6f
--- /dev/null
+++ b/ctdb/tests/simple/55_ctdb_ptrans.sh
@@ -0,0 +1,127 @@
+#!/bin/bash
+
+test_info()
+{
+    cat <<EOF
+Verify that the ctdb ptrans works as expected
+
+Prerequisites:
+
+* An active CTDB cluster with at least 2 active nodes.
+
+Steps:
+
+1. Verify that the status on all of the ctdb nodes is 'OK'.
+2. Pipe some operation to ctdb ptrans and validate the TDB contents with ctdb catdb
+
+Expected results:
+
+* ctdb ptrans works as expected.
+EOF
+}
+
+. "${TEST_SCRIPTS_DIR}/integration.bash"
+
+ctdb_test_init "$@"
+
+set -e
+
+cluster_is_healthy
+
+TESTDB="ptrans_test.tdb"
+
+# Create a temporary persistent database to test with
+echo "create persistent test database $TESTDB"
+try_command_on_node 0 $CTDB attach $TESTDB persistent
+
+# Wipe Test database
+echo "wipe test database"
+try_command_on_node 0 $CTDB wipedb $TESTDB
+
+##########
+
+echo "Adding 3 records"
+
+items='\
+"key1" "value1"
+"key2" "value1"
+"key3" "value1"'
+
+echo "$items" | try_command_on_node -i 0 $CTDB ptrans "$TESTDB"
+
+try_command_on_node 0 $CTDB catdb "$TESTDB"
+
+n=$(echo "$out" | grep -c '^key.*= "key.*"' || true)
+
+if [ $n -ne 3 ] ; then
+    echo "BAD: expected 3 keys in..."
+    echo "$out"
+    exit 1
+else
+    echo "GOOD: 3 records were inserted"
+fi
+
+##########
+
+echo "Deleting 1 record, updating 1, adding 1 new record, 1 bogus input line"
+
+items='\
+"key1" ""
+"key2" "value2"
+"key3"
+"key4" "value1"'
+
+echo "$items" | try_command_on_node -i 0 $CTDB ptrans "$TESTDB"
+
+try_command_on_node 0 $CTDB catdb "$TESTDB"
+
+n=$(echo "$out" | grep -c '^key.*= "key.*"' || true)
+
+if [ $n -ne 3 ] ; then
+    echo "BAD: expected 3 keys in..."
+    echo "$out"
+    exit 1
+else
+    echo "GOOD: 3 records found"
+fi
+
+##########
+
+echo "Verifying records"
+
+while read key value ; do
+    try_command_on_node 0 $CTDB pfetch "$TESTDB" "$key"
+    if [ "$value" != "$out" ] ; then
+	echo "BAD: for key \"$key\" expected \"$value\" but got \"$out\""
+	exit 1
+    else
+	echo "GOOD: for key \"$key\" got \"$out\""
+    fi
+done <<EOF
+key2 value2
+key3 value1
+key4 value1
+EOF
+
+##########
+
+echo "Deleting all records"
+
+items='\
+"key2" ""
+"key3" ""
+"key4" ""'
+
+echo "$items" | try_command_on_node -i 0 $CTDB ptrans "$TESTDB"
+
+try_command_on_node 0 $CTDB catdb "$TESTDB"
+
+n=$(echo "$out" | grep -c '^key.*= "key.*"' || true)
+
+if [ $n -ne 0 ] ; then
+    echo "BAD: expected 0 keys in..."
+    echo "$out"
+    exit 1
+else
+    echo "GOOD: 0 records found"
+fi
diff --git a/ctdb/tools/ctdb.c b/ctdb/tools/ctdb.c
index e906a6939cf..b947705fb28 100644
--- a/ctdb/tools/ctdb.c
+++ b/ctdb/tools/ctdb.c
@@ -4254,6 +4254,155 @@ static int control_pdelete(struct ctdb_context *ctdb, int argc, const char **arg
 	return 0;
 }
 
+static const char *ptrans_parse_string(TALLOC_CTX *mem_ctx, const char *s,
+				       TDB_DATA *data)
+{
+	const char *t;
+	size_t n;
+	const char *ret; /* Next byte after successfully parsed value */
+
+	/* Error, unless someone says otherwise */
+	ret = NULL;
+	/* Indicates no value to parse */
+	*data = tdb_null;
+
+	/* Skip whitespace */
+	n = strspn(s, " \t");
+	t = s + n;
+
+	if (t[0] == '"') {
+		/* Quoted ASCII string - no wide characters! */
+		t++;
+		n = strcspn(t, "\"");
+		if (t[n] == '"') {
+			if (n > 0) {
+				data->dsize = n;
+				data->dptr = talloc_memdup(mem_ctx, t, n);
+				CTDB_NOMEM_ABORT(data->dptr);
+			}
+			ret = t + n + 1;
+		} else {
+			DEBUG(DEBUG_WARNING,("Unmatched \" in input %s\n", s));
+		}
+	} else {
+		DEBUG(DEBUG_WARNING,("Unsupported input format in %s\n", s));
+	}
+
+	return ret;
+}
+
+static bool ptrans_get_key_value(TALLOC_CTX *mem_ctx, FILE *file,
+				 TDB_DATA *key, TDB_DATA *value)
+{
+	char line [1024]; /* FIXME: make this more flexible? */
+	const char *t;
+	char *ptr;
+
+	ptr = fgets(line, sizeof(line), file);
+
+	if (ptr == NULL) {
+		return false;
+	}
+
+	/* Get key */
+	t = ptrans_parse_string(mem_ctx, line, key);
+	if (t == NULL || key->dptr == NULL) {
+		/* Line Ignored but not EOF */
+		return true;
+	}
+
+	/* Get value */
+	t = ptrans_parse_string(mem_ctx, t, value);
+	if (t == NULL) {
+		/* Line Ignored but not EOF */
+		talloc_free(key->dptr);
+		*key = tdb_null;
+		return true;
+	}
+
+	return true;
+}
+
+/*
+ * Update a persistent database as per file/stdin
+ */
+static int control_ptrans(struct ctdb_context *ctdb,
+			  int argc, const char **argv)
+{
+	const char *db_name;
+	struct ctdb_db_context *ctdb_db;
+	TALLOC_CTX *tmp_ctx = talloc_new(ctdb);
+	struct ctdb_transaction_handle *h;
+	TDB_DATA key, value;
+	FILE *file;
+	int ret;
+
+	if (argc != 2) {
+		talloc_free(tmp_ctx);
+		usage();
+	}
+
+	file = stdin;
+	if (strcmp(argv[1], "-") != 0) {
+		file = fopen(argv[1], "r");
+		if (file == NULL) {
+			DEBUG(DEBUG_ERR,("Unable to open file for reading '%s'\n", argv[1]));
+			talloc_free(tmp_ctx);
+			return -1;
+		}
+	}
+
+	db_name = argv[0];
+
+	ctdb_db = ctdb_attach(ctdb, TIMELIMIT(), db_name, true, 0);
+	if (ctdb_db == NULL) {
+		DEBUG(DEBUG_ERR,("Unable to attach to database '%s'\n", db_name));
+		goto error;
+	}
+
+	h = ctdb_transaction_start(ctdb_db, tmp_ctx);
+	if (h == NULL) {
+		DEBUG(DEBUG_ERR,("Failed to start transaction on database %s\n", db_name));
+		goto error;
+	}
+
+	while (ptrans_get_key_value(tmp_ctx, file, &key, &value)) {
+		if (key.dsize != 0) {
+			ret = ctdb_transaction_store(h, key, value);
+			/* Minimise memory use */
+			talloc_free(key.dptr);
+			if (value.dptr != NULL) {
+				talloc_free(value.dptr);
+			}
+			if (ret != 0) {
+				DEBUG(DEBUG_ERR,("Failed to store record\n"));
+				ctdb_transaction_cancel(h);
+				goto error;
+			}
+		}
+	}
+
+	ret = ctdb_transaction_commit(h);
+	if (ret != 0) {
+		DEBUG(DEBUG_ERR,("Failed to commit transaction\n"));
+		goto error;
+	}
+
+	if (file != stdin) {
+		fclose(file);
+	}
+	talloc_free(tmp_ctx);
+	return 0;
+
+error:
+	if (file != stdin) {
+		fclose(file);
+	}
+
+	talloc_free(tmp_ctx);
+	return -1;
+}
+
 /*
   check if a service is bound to a port or not
  */
@@ -6115,6 +6264,7 @@ static const struct {
 	{ "pfetch", 	     control_pfetch,      	false,	false,  "fetch a record from a persistent database", "<dbname|dbid> <key> [<file>]" },
 	{ "pstore", 	     control_pstore,      	false,	false,  "write a record to a persistent database", "<dbname|dbid> <key> <file containing record>" },
 	{ "pdelete", 	     control_pdelete,      	false,	false,  "delete a record from a persistent database", "<dbname|dbid> <key>" },
+	{ "ptrans", 	     control_ptrans,      	false,	false,  "update a persistent database (from stdin)", "<dbname|dbid>" },
 	{ "tfetch", 	     control_tfetch,      	false,	true,  "fetch a record from a [c]tdb-file [-v]", "<tdb-file> <key> [<file>]" },
 	{ "tstore", 	     control_tstore,      	false,	true,  "store a record (including ltdb header)", "<tdb-file> <key> <data+header>" },
 	{ "readkey", 	     control_readkey,      	true,	false,  "read the content off a database key", "<tdb-file> <key>" },