1 files changed, 105 insertions, 0 deletions
diff --git a/clover/collector/grpc/collector_client.py b/clover/collector/grpc/collector_client.py
new file mode 100644
index 0000000..b9e9f67
--- /dev/null
+++ b/clover/collector/grpc/collector_client.py
@@ -0,0 +1,105 @@
+# Copyright (c) Authors of Clover
+#
+# All rights reserved. This program and the accompanying materials
+# are made available under the terms of the Apache License, Version 2.0
+# which accompanies this distribution, and is available at
+# http://www.apache.org/licenses/LICENSE-2.0
+
+from __future__ import print_function
+from kubernetes import client, config
+
+import grpc
+import argparse
+import pickle
+
+import collector_pb2
+import collector_pb2_grpc
+
+# This is a basic client script to test server GRPC messaging
+# TODO improve interface overall
+
+
+def run(args, grpc_port='50054'):
+    pod_ip = get_podip('clover-collector')
+    if pod_ip == '':
+        return "Can not find service: {}".format(args['service_name'])
+    collector_grpc = pod_ip + ':' + grpc_port
+    channel = grpc.insecure_channel(collector_grpc)
+    stub = collector_pb2_grpc.ControllerStub(channel)
+
+    if args['cmd'] == 'init':
+        return init_visibility(stub)
+    elif args['cmd'] == 'start':
+        return start_collector(stub)
+    elif args['cmd'] == 'stop':
+        return stop_collector(stub)
+    elif args['cmd'] == 'clean':
+        return clean_visibility(stub)
+    else:
+        return "Invalid command: {}".format(args['cmd'])
+
+
+def get_podip(pod_name):
+    ip = ''
+    if pod_name != '':
+        config.load_kube_config()
+        v1 = client.CoreV1Api()
+        ret = v1.list_pod_for_all_namespaces(watch=False)
+        for i in ret.items:
+            if i.metadata.name.lower().find(pod_name.lower()) != -1:
+                print("Pod IP: {}".format(i.status.pod_ip))
+                ip = i.status.pod_ip
+                return str(ip)
+    return str(ip)
+
+
+def init_visibility(stub):
+    try:
+        cassandra_hosts = pickle.dumps(['cassandra.default'])
+        response = stub.InitVisibility(collector_pb2.ConfigCassandra(
+            cassandra_hosts=cassandra_hosts, cassandra_port=9042))
+    except Exception as e:
+        return e
+    return response.message
+
+
+def clean_visibility(stub):
+    try:
+        cassandra_hosts = pickle.dumps(['cassandra.default'])
+        schemas = pickle.dumps(['spans', 'traces', 'metrics'])
+        response = stub.TruncateVisibility(collector_pb2.Schemas(
+            schemas=schemas, cassandra_hosts=cassandra_hosts,
+            cassandra_port=9042))
+    except Exception as e:
+        return e
+    return response.message
+
+
+def start_collector(stub):
+    try:
+        cassandra_hosts = pickle.dumps(['cassandra.default'])
+        response = stub.StartCollector(collector_pb2.ConfigCollector(
+            t_port='16686', t_host='jaeger-deployment.istio-system',
+            m_port='9090', m_host='prometheus.istio-system',
+            c_port='9042', c_hosts=cassandra_hosts,
+            sinterval='5'))
+    except Exception as e:
+        return e
+    return response.message
+
+
+def stop_collector(stub):
+    try:
+        response = stub.StopCollector(collector_pb2.ConfigCollector())
+    except Exception as e:
+        return e
+    return response.message
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+            '--cmd', required=True,
+            help='Command to execute in collector')
+    args = parser.parse_args()
+    print(run(vars(args)))