From c6a08f7832bc49076fbe57b1f018be0d47e7f1bd Mon Sep 17 00:00:00 2001
From: Lawrence Lee <lfqlee314@gmail.com>
Date: Mon, 17 Jun 2019 17:24:55 -0700
Subject: [PATCH] [scripts]: add support to db_migrator for non-default unix
 socket (#551)

* [scripts]: add support to db_migrator for non-default unix socket

* add '-s' option to db_migrator

Signed-off-by: Lawrence Lee <t-lale@microsoft.com>

* [scripts]: make db_migrator follow python convention

* change variable names

Signed-off-by: Lawrence Lee <t-lale@microsoft.com>
---
 scripts/db_migrator.py | 28 ++++++++++++++++++++++++----
 1 file changed, 24 insertions(+), 4 deletions(-)

diff --git a/scripts/db_migrator.py b/scripts/db_migrator.py
index effbce0b593a..0eca7c621439 100755
--- a/scripts/db_migrator.py
+++ b/scripts/db_migrator.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python
 
+import traceback
 import sys
 import argparse
 import syslog
@@ -22,7 +23,7 @@ def log_error(msg):
 
 
 class DBMigrator():
-    def __init__(self):
+    def __init__(self, socket=None):
         """
         Version string format:
            version_<major>_<minor>_<build>
@@ -39,7 +40,12 @@ def __init__(self):
         self.TABLE_NAME      = 'VERSIONS'
         self.TABLE_KEY       = 'DATABASE'
         self.TABLE_FIELD     = 'VERSION'
-        self.configDB        = ConfigDBConnector()
+
+        db_kwargs = {}
+        if socket:
+            db_kwargs['unix_socket_path'] = socket
+
+        self.configDB        = ConfigDBConnector(**db_kwargs)
         self.configDB.db_connect('CONFIG_DB')
 
 
@@ -120,16 +126,30 @@ def main():
                             choices=['migrate', 'set_version', 'get_version'],
                             help = 'operation to perform [default: get_version]',
                             default='get_version')
+        parser.add_argument('-s',
+                        dest='socket',
+                        metavar='unix socket',
+                        type = str,
+                        required = False,
+                        help = 'the unix socket that the desired database listens on',
+                        default = None )
         args = parser.parse_args()
         operation = args.operation
+        socket_path = args.socket
+
+        if socket_path:
+            dbmgtr = DBMigrator(socket=socket_path)
+        else:
+            dbmgtr = DBMigrator()
 
-        dbmgtr = DBMigrator()
         result = getattr(dbmgtr, operation)()
         if result:
             print(str(result))
 
     except Exception as e:
-        log_error('Caught excetion: ' + str(e))
+        log_error('Caught exception: ' + str(e))
+        traceback.print_exc()
+        print(str(e))
         parser.print_help()
         sys.exit(1)