myloader.c

/*
    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.

    Authors:        Andrew Hutchings, SkySQL (andrew at skysql dot com)
*/

#define _LARGEFILE64_SOURCE
#define _FILE_OFFSET_BITS 64

#include <mysql.h>

#if defined MARIADB_CLIENT_VERSION_STR && !defined MYSQL_SERVER_VERSION
#define MYSQL_SERVER_VERSION MARIADB_CLIENT_VERSION_STR
#endif

#include <unistd.h>
#include <stdio.h>
#include <string.h>
#include <glib.h>
#include <glib/gstdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <errno.h>
#include <zlib.h>
#include "common.h"
#include "myloader.h"
#include "connection.h"
#include "config.h"
#include "getPassword.h"

guint commit_count = 1000;
gchar *directory = NULL;
gboolean overwrite_tables = FALSE;
gboolean enable_binlog = FALSE;
gchar *source_db = NULL;
static GMutex *init_mutex = NULL;

guint errors = 0;

gboolean read_data(FILE *file, gboolean is_compressed, GString *data,
                   gboolean *eof);
void restore_data(MYSQL *conn, char *database, char *table,
                  const char *filename, gboolean is_schema, gboolean need_use);
void *process_queue(struct thread_data *td);
void add_table(const gchar *filename, struct configuration *conf);
void add_schema(const gchar *filename, MYSQL *conn);
void restore_databases(struct configuration *conf, MYSQL *conn);
void restore_schema_view(MYSQL *conn);
void restore_schema_triggers(MYSQL *conn);
void restore_schema_post(MYSQL *conn);
void no_log(const gchar *log_domain, GLogLevelFlags log_level,
            const gchar *message, gpointer user_data);
void set_verbose(guint verbosity);
void create_database(MYSQL *conn, gchar *database);

static GOptionEntry entries[] = {
    {"directory", 'd', 0, G_OPTION_ARG_STRING, &directory,
     "Directory of the dump to import", NULL},
    {"queries-per-transaction", 'q', 0, G_OPTION_ARG_INT, &commit_count,
     "Number of queries per transaction, default 1000", NULL},
    {"overwrite-tables", 'o', 0, G_OPTION_ARG_NONE, &overwrite_tables,
     "Drop tables if they already exist", NULL},
    {"database", 'B', 0, G_OPTION_ARG_STRING, &db,
     "An alternative database to restore into", NULL},
    {"source-db", 's', 0, G_OPTION_ARG_STRING, &source_db,
     "Database to restore", NULL},
    {"enable-binlog", 'e', 0, G_OPTION_ARG_NONE, &enable_binlog,
     "Enable binary logging of the restore data", NULL},
    {NULL, 0, 0, G_OPTION_ARG_NONE, NULL, NULL, NULL}};

void no_log(const gchar *log_domain, GLogLevelFlags log_level,
            const gchar *message, gpointer user_data) {
  (void)log_domain;
  (void)log_level;
  (void)message;
  (void)user_data;
}

void set_verbose(guint verbosity) {
  switch (verbosity) {
  case 0:
    g_log_set_handler(NULL, (GLogLevelFlags)(G_LOG_LEVEL_MASK), no_log, NULL);
    break;
  case 1:
    g_log_set_handler(
        NULL, (GLogLevelFlags)(G_LOG_LEVEL_WARNING | G_LOG_LEVEL_MESSAGE),
        no_log, NULL);
    break;
  case 2:
    g_log_set_handler(NULL, (GLogLevelFlags)(G_LOG_LEVEL_MESSAGE), no_log,
                      NULL);
    break;
  default:
    break;
  }
}

int main(int argc, char *argv[]) {
  struct configuration conf = {NULL, NULL, NULL, 0};

  GError *error = NULL;
  GOptionContext *context;

  g_thread_init(NULL);

  init_mutex = g_mutex_new();

  if (db == NULL && source_db != NULL) {
    db = g_strdup(source_db);
  }

  context = g_option_context_new("multi-threaded MySQL loader");
  GOptionGroup *main_group =
      g_option_group_new("main", "Main Options", "Main Options", NULL, NULL);
  g_option_group_add_entries(main_group, entries);
  g_option_group_add_entries(main_group, common_entries);
  g_option_context_set_main_group(context, main_group);
  if (!g_option_context_parse(context, &argc, &argv, &error)) {
    g_print("option parsing failed: %s, try --help\n", error->message);
    exit(EXIT_FAILURE);
  }
  g_option_context_free(context);

  // prompt for password if it's NULL
  if (sizeof(password) == 0 || (password == NULL && askPassword)) {
    password = passwordPrompt();
  }

  if (program_version) {
    g_print("myloader %s, built against MySQL %s\n", VERSION,
            MYSQL_VERSION_STR);
    exit(EXIT_SUCCESS);
  }

  set_verbose(verbose);

  if (!directory) {
    g_critical("a directory needs to be specified, see --help\n");
    exit(EXIT_FAILURE);
  } else {
    char *p = g_strdup_printf("%s/metadata", directory);
    if (!g_file_test(p, G_FILE_TEST_EXISTS)) {
      g_critical("the specified directory is not a mydumper backup\n");
      exit(EXIT_FAILURE);
    }
  }
  MYSQL *conn;
  conn = mysql_init(NULL);

  configure_connection(conn, "myloader");
  if (!mysql_real_connect(conn, hostname, username, password, NULL, port,
                          socket_path, 0)) {
    g_critical("Error connection to database: %s", mysql_error(conn));
    exit(EXIT_FAILURE);
  }

  if (mysql_query(conn, "SET SESSION wait_timeout = 2147483")) {
    g_warning("Failed to increase wait_timeout: %s", mysql_error(conn));
  }

  if (!enable_binlog)
    mysql_query(conn, "SET SQL_LOG_BIN=0");

  mysql_query(conn, "/*!40014 SET FOREIGN_KEY_CHECKS=0*/");
  conf.queue = g_async_queue_new();
  conf.ready = g_async_queue_new();

  guint n;
  GThread **threads = g_new(GThread *, num_threads);
  struct thread_data *td = g_new(struct thread_data, num_threads);
  for (n = 0; n < num_threads; n++) {
    td[n].conf = &conf;
    td[n].thread_id = n + 1;
    threads[n] =
        g_thread_create((GThreadFunc)process_queue, &td[n], TRUE, NULL);
    g_async_queue_pop(conf.ready);
  }
  g_async_queue_unref(conf.ready);

  g_message("%d threads created", num_threads);

  restore_databases(&conf, conn);

  for (n = 0; n < num_threads; n++) {
    struct job *j = g_new0(struct job, 1);
    j->type = JOB_SHUTDOWN;
    g_async_queue_push(conf.queue, j);
  }

  for (n = 0; n < num_threads; n++) {
    g_thread_join(threads[n]);
  }

  restore_schema_post(conn);

  restore_schema_view(conn);

  restore_schema_triggers(conn);

  g_async_queue_unref(conf.queue);
  mysql_close(conn);
  mysql_thread_end();
  mysql_library_end();
  g_free(directory);
  g_free(td);
  g_free(threads);

  return errors ? EXIT_FAILURE : EXIT_SUCCESS;
}

void restore_databases(struct configuration *conf, MYSQL *conn) {
  GError *error = NULL;
  GDir *dir = g_dir_open(directory, 0, &error);

  if (error) {
    g_critical("cannot open directory %s, %s\n", directory, error->message);
    errors++;
    return;
  }

  const gchar *filename = NULL;

  while ((filename = g_dir_read_name(dir))) {
    if (!source_db ||
        g_str_has_prefix(filename, g_strdup_printf("%s.", source_db))) {
      if (g_strrstr(filename, "-schema.sql")) {
        add_schema(filename, conn);
      }
    }
  }

  g_dir_rewind(dir);

  while ((filename = g_dir_read_name(dir))) {
    if (!source_db ||
        g_str_has_prefix(filename, g_strdup_printf("%s.", source_db))) {
      if (!g_strrstr(filename, "-schema.sql") &&
          !g_strrstr(filename, "-schema-view.sql") &&
          !g_strrstr(filename, "-schema-triggers.sql") &&
          !g_strrstr(filename, "-schema-post.sql") &&
          !g_strrstr(filename, "-schema-create.sql") &&
          g_strrstr(filename, ".sql")) {
        add_table(filename, conf);
      }
    }
  }

  g_dir_close(dir);
}

void restore_schema_view(MYSQL *conn) {
  GError *error = NULL;
  GDir *dir = g_dir_open(directory, 0, &error);

  if (error) {
    g_critical("cannot open directory %s, %s\n", directory, error->message);
    errors++;
    return;
  }

  const gchar *filename = NULL;

  while ((filename = g_dir_read_name(dir))) {
    if (!source_db ||
        g_str_has_prefix(filename, g_strdup_printf("%s.", source_db))) {
      if (g_strrstr(filename, "-schema-view.sql")) {
        add_schema(filename, conn);
      }
    }
  }

  g_dir_close(dir);
}

void restore_schema_triggers(MYSQL *conn) {
  GError *error = NULL;
  GDir *dir = g_dir_open(directory, 0, &error);
  gchar **split_file = NULL;
  gchar *database = NULL;
  gchar **split_table = NULL;
  gchar *table = NULL;

  if (error) {
    g_critical("cannot open directory %s, %s\n", directory, error->message);
    errors++;
    return;
  }

  const gchar *filename = NULL;

  while ((filename = g_dir_read_name(dir))) {
    if (!source_db ||
        g_str_has_prefix(filename, g_strdup_printf("%s.", source_db))) {
      if (g_strrstr(filename, "-schema-triggers.sql")) {
        split_file = g_strsplit(filename, ".", 0);
        database = split_file[0];
        split_table = g_strsplit(split_file[1], "-schema", 0);
        table = split_table[0];
        g_message("Restoring triggers for `%s`.`%s`", db ? db : database,
                  table);
        restore_data(conn, database, table, filename, TRUE, TRUE);
      }
    }
  }

  g_strfreev(split_table);
  g_strfreev(split_file);
  g_dir_close(dir);
}

void restore_schema_post(MYSQL *conn) {
  GError *error = NULL;
  GDir *dir = g_dir_open(directory, 0, &error);
  gchar **split_file = NULL;
  gchar *database = NULL;
  // gchar* table=NULL;

  if (error) {
    g_critical("cannot open directory %s, %s\n", directory, error->message);
    errors++;
    return;
  }

  const gchar *filename = NULL;

  while ((filename = g_dir_read_name(dir))) {
    if (!source_db ||
        g_str_has_prefix(filename, g_strdup_printf("%s.", source_db))) {
      if (g_strrstr(filename, "-schema-post.sql")) {
        split_file = g_strsplit(filename, "-schema-post.sql", 0);
        database = split_file[0];
        // table= split_file[0]; //NULL
        g_message("Restoring routines and events for `%s`", db ? db : database);
        restore_data(conn, database, NULL, filename, TRUE, TRUE);
      }
    }
  }

  g_strfreev(split_file);
  g_dir_close(dir);
}

void create_database(MYSQL *conn, gchar *database) {

  gchar *query = NULL;

  if ((db == NULL && source_db == NULL) ||
      (db != NULL && source_db != NULL && !g_ascii_strcasecmp(db, source_db))) {
    const gchar *filename =
        g_strdup_printf("%s-schema-create.sql", db ? db : database);
    const gchar *filenamegz =
        g_strdup_printf("%s-schema-create.sql.gz", db ? db : database);
    const gchar *filepath = g_strdup_printf("%s/%s-schema-create.sql",
                                            directory, db ? db : database);
    const gchar *filepathgz = g_strdup_printf("%s/%s-schema-create.sql.gz",
                                              directory, db ? db : database);

    if (g_file_test(filepath, G_FILE_TEST_EXISTS)) {
      restore_data(conn, database, NULL, filename, TRUE, FALSE);
    } else if (g_file_test(filepathgz, G_FILE_TEST_EXISTS)) {
      restore_data(conn, database, NULL, filenamegz, TRUE, FALSE);
    } else {
      query = g_strdup_printf("CREATE DATABASE `%s`", db ? db : database);
      mysql_query(conn, query);
    }
  } else {
    query = g_strdup_printf("CREATE DATABASE `%s`", db ? db : database);
    mysql_query(conn, query);
  }

  g_free(query);
  return;
}

void add_schema(const gchar *filename, MYSQL *conn) {
  // 0 is database, 1 is table with -schema on the end
  gchar **split_file = g_strsplit(filename, ".", 0);
  gchar *database = split_file[0];
  // Remove the -schema from the table name
  gchar **split_table = g_strsplit(split_file[1], "-schema", 0);
  gchar *table = split_table[0];

  gchar *query =
      g_strdup_printf("SHOW CREATE DATABASE `%s`", db ? db : database);
  if (mysql_query(conn, query)) {
    g_message("Creating database `%s`", db ? db : database);
    create_database(conn, database);
  } else {
    MYSQL_RES *result = mysql_store_result(conn);
    // In drizzle the query succeeds with no rows
    my_ulonglong row_count = mysql_num_rows(result);
    mysql_free_result(result);
    if (row_count == 0) {
      create_database(conn, database);
    }
  }

  if (overwrite_tables) {
    g_message("Dropping table or view (if exists) `%s`.`%s`",
              db ? db : database, table);
    query = g_strdup_printf("DROP TABLE IF EXISTS `%s`.`%s`",
                            db ? db : database, table);
    mysql_query(conn, query);
    query = g_strdup_printf("DROP VIEW IF EXISTS `%s`.`%s`", db ? db : database,
                            table);
    mysql_query(conn, query);
  }

  g_free(query);

  g_message("Creating table `%s`.`%s`", db ? db : database, table);
  restore_data(conn, database, table, filename, TRUE, TRUE);
  g_strfreev(split_table);
  g_strfreev(split_file);
  return;
}

void add_table(const gchar *filename, struct configuration *conf) {
  struct job *j = g_new0(struct job, 1);
  struct restore_job *rj = g_new(struct restore_job, 1);
  j->job_data = (void *)rj;
  rj->filename = g_strdup(filename);
  j->type = JOB_RESTORE;
  gchar **split_file = g_strsplit(filename, ".", 0);
  rj->database = g_strdup(split_file[0]);
  rj->table = g_strdup(split_file[1]);
  rj->part = g_ascii_strtoull(split_file[2], NULL, 10);
  g_async_queue_push(conf->queue, j);
  return;
}

void *process_queue(struct thread_data *td) {
  struct configuration *conf = td->conf;
  g_mutex_lock(init_mutex);
  MYSQL *thrconn = mysql_init(NULL);
  g_mutex_unlock(init_mutex);

  configure_connection(thrconn, "myloader");

  if (!mysql_real_connect(thrconn, hostname, username, password, NULL, port,
                          socket_path, 0)) {
    g_critical("Failed to connect to MySQL server: %s", mysql_error(thrconn));
    exit(EXIT_FAILURE);
  }

  if (mysql_query(thrconn, "SET SESSION wait_timeout = 2147483")) {
    g_warning("Failed to increase wait_timeout: %s", mysql_error(thrconn));
  }

  if (!enable_binlog)
    mysql_query(thrconn, "SET SQL_LOG_BIN=0");

  mysql_query(thrconn, "/*!40101 SET NAMES binary*/");
  mysql_query(thrconn, "/*!40101 SET SQL_MODE='NO_AUTO_VALUE_ON_ZERO' */");
  mysql_query(thrconn, "/*!40014 SET UNIQUE_CHECKS=0 */");
  mysql_query(thrconn, "SET autocommit=0");

  g_async_queue_push(conf->ready, GINT_TO_POINTER(1));

  struct job *job = NULL;
  struct restore_job *rj = NULL;
  for (;;) {
    job = (struct job *)g_async_queue_pop(conf->queue);

    switch (job->type) {
    case JOB_RESTORE:
      rj = (struct restore_job *)job->job_data;
      g_message("Thread %d restoring `%s`.`%s` part %d", td->thread_id,
                rj->database, rj->table, rj->part);
      restore_data(thrconn, rj->database, rj->table, rj->filename, FALSE, TRUE);
      if (rj->database)
        g_free(rj->database);
      if (rj->table)
        g_free(rj->table);
      if (rj->filename)
        g_free(rj->filename);
      g_free(rj);
      g_free(job);
      break;
    case JOB_SHUTDOWN:
      g_message("Thread %d shutting down", td->thread_id);
      if (thrconn)
        mysql_close(thrconn);
      g_free(job);
      mysql_thread_end();
      return NULL;
      break;
    default:
      g_critical("Something very bad happened!");
      exit(EXIT_FAILURE);
    }
  }
  if (thrconn)
    mysql_close(thrconn);
  mysql_thread_end();
  return NULL;
}

void restore_data(MYSQL *conn, char *database, char *table,
                  const char *filename, gboolean is_schema, gboolean need_use) {
  void *infile;
  gboolean is_compressed = FALSE;
  gboolean eof = FALSE;
  guint query_counter = 0;
  GString *data = g_string_sized_new(512);

  gchar *path = g_build_filename(directory, filename, NULL);

  if (!g_str_has_suffix(path, ".gz")) {
    infile = g_fopen(path, "r");
    is_compressed = FALSE;
  } else {
    infile = (void *)gzopen(path, "r");
    is_compressed = TRUE;
  }

  if (!infile) {
    g_critical("cannot open file %s (%d)", filename, errno);
    errors++;
    return;
  }

  if (need_use) {
    gchar *query = g_strdup_printf("USE `%s`", db ? db : database);

    if (mysql_query(conn, query)) {
      g_critical("Error switching to database %s whilst restoring table %s",
                 db ? db : database, table);
      g_free(query);
      errors++;
      return;
    }

    g_free(query);
  }

  if (!is_schema)
    mysql_query(conn, "START TRANSACTION");

  while (eof == FALSE) {
    if (read_data(infile, is_compressed, data, &eof)) {
      // Search for ; in last 5 chars of line
      if (g_strrstr(&data->str[data->len >= 5 ? data->len - 5 : 0], ";\n")) {
        if (mysql_real_query(conn, data->str, data->len)) {
          g_critical("Error restoring %s.%s from file %s: %s",
                     db ? db : database, table, filename, mysql_error(conn));
          errors++;
          return;
        }
        query_counter++;
        if (!is_schema && (query_counter == commit_count)) {
          query_counter = 0;
          if (mysql_query(conn, "COMMIT")) {
            g_critical("Error committing data for %s.%s: %s",
                       db ? db : database, table, mysql_error(conn));
            errors++;
            return;
          }
          mysql_query(conn, "START TRANSACTION");
        }

        g_string_set_size(data, 0);
      }
    } else {
      g_critical("error reading file %s (%d)", filename, errno);
      errors++;
      return;
    }
  }
  if (!is_schema && mysql_query(conn, "COMMIT")) {
    g_critical("Error committing data for %s.%s from file %s: %s",
               db ? db : database, table, filename, mysql_error(conn));
    errors++;
  }
  g_string_free(data, TRUE);
  g_free(path);
  if (!is_compressed) {
    fclose(infile);
  } else {
    gzclose((gzFile)infile);
  }
  return;
}

gboolean read_data(FILE *file, gboolean is_compressed, GString *data,
                   gboolean *eof) {
  char buffer[256];

  do {
    if (!is_compressed) {
      if (fgets(buffer, 256, file) == NULL) {
        if (feof(file)) {
          *eof = TRUE;
          buffer[0] = '\0';
        } else {
          return FALSE;
        }
      }
    } else {
      if (!gzgets((gzFile)file, buffer, 256)) {
        if (gzeof((gzFile)file)) {
          *eof = TRUE;
          buffer[0] = '\0';
        } else {
          return FALSE;
        }
      }
    }
    g_string_append(data, buffer);
  } while ((buffer[strlen(buffer)] != '\0') && *eof == FALSE);

  return TRUE;
}