add commandline flag to specify encoding

nmandery · Feb 28, 2012 · 8bcf494 · 8bcf494
1 parent 49d668c
commit 8bcf494
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 7 deletions.
diff --git a/README.md b/README.md
@@ -25,6 +25,7 @@ Result:
 
     Usage :  
     psqlchunks command [options] files
+    version: 0.5.1
     use - as filename to read from stdin.
 
     Definition of a chunk of SQL:
@@ -68,7 +69,10 @@ Result:
       -a           abort execution after first failed chunk. (default: continue)
       -l           number of lines to output before and after failing lines
                    of SQL. (default: 2)
-
+      -E           set the client_encoding of the database connection. This
+                   setting is usefull when the encoding of sql file differs
+                   from the default client_encoding of the database server.
+
     Connection parameters:
       -d [database name]
       -U [user]

diff --git a/src/db.cc b/src/db.cc
@@ -28,6 +28,28 @@ Db::connect( const char * host, const char * db_name,  const char * port, const
     return isConnected();
 }
 
+
+bool
+Db::setEncoding(const char * enc_name)
+{
+    if (enc_name == NULL) {
+        return false;
+    }
+
+    std::stringstream sqlstrm;
+    sqlstrm << "set client_encoding to " << enc_name << ";";
+    std::string sqlstr = sqlstrm.str();
+
+    try {
+        executeSql(sqlstr.c_str(), true);
+    }
+    catch (DbException &e) {
+        return false;
+    }
+
+    return true;
+}
+
 void
 Db::disconnect()
 {
@@ -155,7 +177,7 @@ Db::finish()
 
 
 void
-Db::executeSql(const char * sqlstr)
+Db::executeSql(const char * sqlstr, bool silent)
 {
     if (!isConnected()) {
         log_warn("can not run chunk - no db connection");
@@ -181,7 +203,9 @@ Db::executeSql(const char * sqlstr)
 
         PQclear(pgres);
 
-        log_error("%s", msg.c_str());
+        if (!silent) {
+            log_error("%s", msg.c_str());
+        }
         DbException e(msg);
         throw e;
     }

diff --git a/src/db.h b/src/db.h
@@ -37,7 +37,11 @@ namespace PsqlChunks
             void commit();
             void rollback();
             void begin();
-            void executeSql(const char *);
+
+            /**
+             * silent: do not log on error
+             */
+            void executeSql(const char *, bool silent = false);
 
 
 
@@ -65,6 +69,8 @@ namespace PsqlChunks
                 return failed_count;
             }
 
+            bool setEncoding(const char * enc_name);
+
             bool runChunk(Chunk & chunk);
             void finish();
             bool cancel(std::string &);

diff --git a/src/psqlchunks.cc b/src/psqlchunks.cc
@@ -44,7 +44,7 @@ using namespace PsqlChunks;
 // version number
 #define VERSION_MAJOR 0
 #define VERSION_MINOR 5
-#define VERSION_PATCH 0
+#define VERSION_PATCH 1
 
 static const char * s_fail_sep = "-------------------------------------------------------";
 
@@ -90,6 +90,7 @@ struct Settings {
     bool is_terminal;
     unsigned int context_lines;
     bool print_filenames;
+    const char * client_encoding;
 };
 static Settings settings = {
     NULL,       /* db_port */
@@ -102,7 +103,8 @@ static Settings settings = {
     LIST,       /* command */
     false,      /* colored */
     DEFAULT_CONTEXT_LINES,   /* context_lines */
-    true        /* print_filenames */
+    true,       /* print_filenames */
+    NULL        /* client_encoding */
 };
 
 
@@ -195,6 +197,9 @@ print_help()
         "  -a           abort execution after first failed chunk. (default: continue)\n"
         "  -l           number of lines to output before and after failing lines\n"
         "               of SQL. (default: " STRINGIFY(DEFAULT_CONTEXT_LINES) ")\n"
+        "  -E           set the client_encoding of the database connection. This\n"
+        "               setting is usefull when the encoding of sql file differs\n"
+        "               from the default client_encoding of the database server.\n"
         "\n"
         "Connection parameters:\n"
         "  -d [database name]\n"
@@ -439,6 +444,15 @@ handle_files(char * files[], int nufiles)
                 fprintf(stderr, "%s\n", db.getErrorMessage().c_str());
                 rc = RC_E_USAGE;
             }
+            else {
+                if (settings.client_encoding != NULL) {
+                    if (!db.setEncoding(settings.client_encoding)) {
+                        fprintf(stderr, "Could not set encoding to %s.\nPostgresql message: %s\n", settings.client_encoding,
+                                    db.getErrorMessage().c_str());
+                        rc = RC_E_USAGE;
+                    }
+                }
+            }
 
             db.setCommit(settings.commit_sql);
         }
@@ -519,7 +533,7 @@ main(int argc, char * argv[] )
 
     // read options
     char opt;
-    while ( (opt = getopt(argc, argv, "l:p:U:d:h:WCaF")) != -1) {
+    while ( (opt = getopt(argc, argv, "l:p:U:d:h:WCaFE:")) != -1) {
         switch (opt) {
             case 'p': /* port */
                 settings.db_port = optarg;
@@ -562,6 +576,9 @@ main(int argc, char * argv[] )
             case 'F':
                 settings.print_filenames = false;
                 break;
+            case 'E': /* client_encoding */
+                settings.client_encoding = optarg;
+                break;
             default:
                 quit("Unknown option.");
         }