PySpark integration with client (mc2-project#261)

This PR includes the necessary changes to make our PySpark listener actually work with the full end-to-end workflow.

src/python/intp_handler.py

@@ -1,19 +1,24 @@
-import code
-import io
-from code import InteractiveInterpreter
 from contextlib import redirect_stdout, redirect_stderr
+from io import StringIO
+from traceback import print_exc
 
 from opaque_sql import *
 
+from pyspark.shell import *
+
 class IntpHandler:
     def __init__(self):
-        init_opaque_sql()
+        self.initialized = False
 
     def run(self, source):
-        out = io.StringIO()
-        err = io.StringIO()
-        with redirect_stdout(out) and redirect_stderr(err):
-            compiled = compile(source)
-            exec(compiled)
-        out.getvalue(), err.getvalue()
-
+        with StringIO() as out, redirect_stdout(out), \
+                StringIO() as err, redirect_stderr(err):
+            try:
+                if not self.initialized:
+                    exec("init_opaque_sql()")
+                    self.initialized = True
+                exec(source)
+            except Exception as e:
+                print_exc() # This goes to stderr
+            finally:
+                return out.getvalue(), err.getvalue()

src/python/listener.py

@@ -11,11 +11,11 @@ def __init__(self):
 
     def ReceiveQuery(self, request, context):
         source = request.request
-        output, result = self.intp_handler.run(source)
-        status = listener_pb2.Status(0, "")
-        if result: # code threw an exception
-            status = listener_pb2.Status(1, result)
-        return listener_pb2.QueryResult(output, status)
+        out, err = self.intp_handler.run(source)
+        status = opaquesql_pb2.Status(status=0, exception="")
+        if err: # If non-empty, user code threw an exception
+            status = opaquesql_pb2.Status(status=1, exception="Opaque SQL returned an error: " + err)
+        return opaquesql_pb2.QueryResult(result=out, status=status)
 
 if __name__ == "__main__":
     server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))