Merge pull request #26 from scalableminds/grpc-health-in-client-proposal

refactor python client

Author: fm3

client/fossildb-client

@@ -10,16 +10,8 @@ import fossildbapi_pb2_grpc as proto_rpc
 from grpc_health.v1 import health_pb2
 from grpc_health.v1 import health_pb2_grpc
 
-def main():
-
-    commands = {
-        'backup': lambda stub:
-            stub.Backup(proto.BackupRequest()),
-        'restore': lambda stub:
-            stub.RestoreFromBackup(proto.RestoreFromBackupRequest()),
-        'health': None
-    }
 
+def parse_args(commands):
     parser = argparse.ArgumentParser()
     parser.add_argument(
         'address', metavar='address', default='localhost', nargs='?',
@@ -32,39 +24,48 @@ def main():
         help='command to execute, one of {}'.format(list(commands.keys())))
 
     args = parser.parse_args()
-
     if args.command not in commands:
         print("command {} is not available".format(args.command))
         parser.print_help()
         exit(20)
 
+    return args
+
+
+def health(channel):
+    try :
+        healthStub = health_pb2_grpc.HealthStub(channel)
+        reply = healthStub.Check(health_pb2.HealthCheckRequest(service=''))
+        STATUSMAP = health_pb2._HEALTHCHECKRESPONSE_SERVINGSTATUS.values_by_name
+        SERVING = STATUSMAP["SERVING"].number
+        if reply.status != SERVING:
+            raise Exception(reply.status)
+    except Exception as e:
+        print('Health check unsuccessful. FossilDB offline?')
+        print(e)
+        sys.exit(1)
+    return reply
+
+
+def main():
+    commands = {
+        'backup': lambda channel:
+            proto_rpc.FossilDBStub(channel).Backup(proto.BackupRequest()),
+        'restore': lambda channel:
+            proto_rpc.FossilDBStub(channel).RestoreFromBackup(proto.RestoreFromBackupRequest()),
+        'health': health
+    }
+
+    args = parse_args(commands)
     full_address = '{}:{}'.format(args.address, args.port)
 
     print('Connecting to FossilDB at', full_address)
-
     channel = grpc.insecure_channel(full_address)
-    stub = proto_rpc.FossilDBStub(channel)
-    healthStub = health_pb2_grpc.HealthStub(channel)
-
-    if args.command == 'health':
-        try :
-            reply = healthStub.Check(health_pb2.HealthCheckRequest(service=''))
-            print(reply)
-            UNKNOWN = 0;
-            SERVING = 1;
-            NOT_SERVING = 2;
-            if reply.status != SERVING:
-                raise Exception(reply.status)
-        except Exception as e:
-            print('Health check unsuccessful. FossilDB offline?')
-            print(e)
-            sys.exit(1)
-
-    else:
-        reply = commands[args.command](stub)
-        print(reply)
-        if not reply.success:
-            sys.exit(1)
+
+    reply = commands[args.command](channel)
+    print(reply)
+    if hasattr(reply, 'success') and not reply.success:
+        sys.exit(1)
 
 
 if __name__ == '__main__':