diff options
Diffstat (limited to 'examples/mongo-dump.c')
| -rw-r--r-- | examples/mongo-dump.c | 224 |
1 files changed, 224 insertions, 0 deletions
diff --git a/examples/mongo-dump.c b/examples/mongo-dump.c new file mode 100644 index 0000000..7e6419f --- /dev/null +++ b/examples/mongo-dump.c @@ -0,0 +1,224 @@ +/* mongo-dump.c - MongoDB database dumper; example application. + * Copyright 2011, 2012 Gergely Nagy <algernon@balabit.hu> + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include <mongo.h> + +#include <unistd.h> +#include <stdio.h> +#include <stdlib.h> +#include <errno.h> +#include <string.h> +#include <sys/types.h> +#include <sys/stat.h> +#include <fcntl.h> + +typedef struct +{ + gchar *addr; + gint port; + gchar *db; + gchar *coll; + gchar *output; + gchar *ns; + gboolean verbose; + gboolean slaveok; + gboolean master_sync; +} config_t; + +#define VLOG(...) { if (config->verbose) fprintf (stderr, __VA_ARGS__); } + +void +mongo_dump (config_t *config) +{ + mongo_sync_connection *conn; + mongo_sync_cursor *cursor; + bson *b; + int fd; + + glong cnt, pos = 0; + + gchar *error = NULL; + int e; + + if (config->port == MONGO_CONN_LOCAL) + { + VLOG ("Connecting to %s/%s.%s...\n", config->addr, config->db, + config->coll); + } + else + { + VLOG ("Connecting to %s:%d/%s.%s...\n", config->addr, config->port, + config->db, config->coll); + } + conn = mongo_sync_connect (config->addr, config->port, config->slaveok); + + if (!conn) + { + e = errno; + + mongo_sync_cmd_get_last_error (conn, config->db, &error); + fprintf (stderr, "Error connecting to %s:%d: %s\n", config->addr, + config->port, (error) ? error : strerror (e)); + g_free (error); + exit (1); + } + + if (config->master_sync) + { + VLOG ("Syncing to master...\n"); + conn = mongo_sync_reconnect (conn, TRUE); + if (!conn) + { + e = errno; + + mongo_sync_cmd_get_last_error (conn, config->db, &error); + fprintf (stderr, "Error reconnecting to the master of %s:%d: %s\n", + config->addr, config->port, (error) ? error : strerror (e)); + exit (1); + } + } + + VLOG ("Counting documents...\n"); + cnt = mongo_sync_cmd_count (conn, config->db, config->coll, NULL); + if (cnt < 0) + { + e = errno; + + mongo_sync_cmd_get_last_error (conn, config->db, &error); + fprintf (stderr, "Error counting documents in %s.%s: %s\n", + config->db, config->coll, (error) ? error : strerror (e)); + mongo_sync_disconnect (conn); + exit (1); + } + + VLOG ("Opening output file '%s'...\n", config->output); + if (strcmp (config->output, "-") == 0) + fd = 1; + else + { + fd = open (config->output, O_RDWR | O_CREAT | O_TRUNC, 0600); + if (fd == -1) + { + fprintf (stderr, "Error opening output file '%s': %s\n", + config->output, strerror (errno)); + mongo_sync_disconnect (conn); + exit (1); + } + } + + VLOG ("Launching initial query...\n"); + b = bson_new (); + bson_finish (b); + cursor = mongo_sync_cursor_new (conn, config->ns, + mongo_sync_cmd_query (conn, config->ns, + MONGO_WIRE_FLAG_QUERY_NO_CURSOR_TIMEOUT, + 0, 10, b, NULL)); + bson_free (b); + + while ((pos < cnt) && mongo_sync_cursor_next (cursor)) + { + bson *b = mongo_sync_cursor_get_data (cursor); + pos++; + + if (!b) + { + e = errno; + + mongo_sync_cmd_get_last_error (conn, config->db, &error); + fprintf (stderr, "Error advancing the cursor: %s\n", + (error) ? error : strerror (e)); + mongo_sync_disconnect (conn); + exit (1); + } + + if (pos % 10 == 0) + VLOG ("\rDumping... %03.2f%%", (pos * 1.0) / (cnt * 1.0) * 100); + + if (write (fd, bson_data (b), bson_size (b)) != bson_size (b)) + { + perror ("write()"); + exit (1); + } + bson_free (b); + } + VLOG ("\rDumping... %03.2f%%\n", (double)((pos / cnt) * 100)); + + mongo_sync_cursor_free (cursor); + + close (fd); + mongo_sync_disconnect (conn); +} + +int +main (int argc, char *argv[]) +{ + GError *error = NULL; + GOptionContext *context; + config_t config = { + NULL, 27017, NULL, NULL, NULL, NULL, FALSE, FALSE, FALSE + }; + + GOptionEntry entries[] = + { + { "addr", 'a', 0, G_OPTION_ARG_STRING, &config.addr, + "Address to connect to", "ADDRESS" }, + { "port", 'p', 0, G_OPTION_ARG_INT, &config.port, "Port", "PORT" }, + { "db", 'd', 0, G_OPTION_ARG_STRING, &config.db, "Database", "DB" }, + { "collection", 'c', 0, G_OPTION_ARG_STRING, &config.coll, "Collection", + "COLL" }, + { "verbose", 'v', 0, G_OPTION_ARG_NONE, &config.verbose, + "Be verbose", NULL }, + { "output", 'o', 0, G_OPTION_ARG_STRING, &config.output, + "Output", "FILENAME" }, + { "slave-ok", 's', 0, G_OPTION_ARG_NONE, &config.slaveok, + "Connecting to slaves is ok", NULL }, + { "master-sync", 'm', 0, G_OPTION_ARG_NONE, &config.master_sync, + "Reconnect to the replica master", NULL }, + { NULL, 0, 0, 0, NULL, NULL, NULL } + }; + + context = g_option_context_new ("- dump a complete mongo collection"); + g_option_context_add_main_entries (context, entries, "mongo-dump"); + if (!g_option_context_parse (context, &argc, &argv, &error)) + { + g_print ("option parsing failed: %s\n", error->message); + exit (1); + } + + if (!((config.addr && config.port)) || !config.db || + !config.coll || !config.output) + { + gchar **nargv; + argc = 2; + + nargv = g_new (gchar *, 3); + nargv[0] = argv[0]; + nargv[1] = "--help"; + nargv[2] = NULL; + + g_option_context_parse (context, &argc, (gchar ***)&nargv, &error); + + exit (1); + } + + config.ns = g_strdup_printf ("%s.%s", config.db, config.coll); + mongo_dump (&config); + + g_free (config.ns); + g_option_context_free (context); + + return 0; +} |