From 62f06c645acda2eee538b211bc8131c9ab30b315 Mon Sep 17 00:00:00 2001
From: Chad Austin <chadaustin@fb.com>
Date: Tue, 15 Sep 2020 11:39:35 -0700
Subject: [PATCH] introduce a py3 client

Summary: Introduce a py3 replacement for EdenClient.

Reviewed By: wez

Differential Revision: D22060644

fbshipit-source-id: 535a1f1addf3ac79de0471a688cef55567802d26
---
 eden/fs/py/eden/thrift/client.py           | 70 ++++++++++++++++++++++
 eden/fs/py/eden/thrift/test/client_test.py | 26 ++++++++
 2 files changed, 96 insertions(+)
 create mode 100644 eden/fs/py/eden/thrift/client.py
 create mode 100644 eden/fs/py/eden/thrift/test/client_test.py

diff --git a/eden/fs/py/eden/thrift/client.py b/eden/fs/py/eden/thrift/client.py
new file mode 100644
index 0000000000..c98df6d43a
--- /dev/null
+++ b/eden/fs/py/eden/thrift/client.py
@@ -0,0 +1,70 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This software may be used and distributed according to the terms of the
+# GNU General Public License version 2.
+
+import os
+from typing import Optional
+
+from eden.fs.service.eden.clients import EdenService
+from thrift.py3 import get_client
+from thrift.py3.exceptions import TransportError, TransportErrorType
+
+
+SOCKET_PATH = "socket"
+
+
+class EdenNotRunningError(Exception):
+    def __init__(self, socket_path: str):
+        super(EdenNotRunningError, self).__init__(
+            f"edenfs daemon does not appear to be running: tried {socket_path}"
+        )
+        self.socket_path = socket_path
+
+
+class EdenClient(EdenService):
+    """
+    EdenClient is a subclass of EdenService that provides
+    some conveniences and helps deal with evolving Thrift APIs.
+    """
+
+    socket_path: Optional[str] = None
+
+    async def __aenter__(self):
+        try:
+            return await super().__aenter__()
+        except TransportError as ex:
+            if ex.type == TransportErrorType.NOT_OPEN:
+                raise EdenNotRunningError(self.socket_path)
+            raise
+
+
+def create_thrift_client(
+    eden_dir: Optional[str] = None,
+    socket_path: Optional[str] = None,
+    timeout: Optional[float] = None,
+) -> EdenClient:
+    """
+    Construct a thrift client to speak to the running eden server
+    instance associated with the specified mount point.
+
+    @return Returns an EdenService.Client.
+    """
+
+    if socket_path is not None:
+        pass
+    elif eden_dir is not None:
+        socket_path = os.path.join(eden_dir, SOCKET_PATH)
+    else:
+        raise TypeError("one of eden_dir or socket_path is required")
+
+    if timeout is None:
+        # We used to set a default timeout here, but picking the right duration is hard,
+        # and safely retrying an arbitrary thrift call may not be safe.  So we
+        # just leave the client with no timeout, unless one is given.
+        timeout = 0
+
+    client = get_client(EdenClient, path=socket_path, timeout=timeout)
+    client.socket_path = socket_path
+    return client
diff --git a/eden/fs/py/eden/thrift/test/client_test.py b/eden/fs/py/eden/thrift/test/client_test.py
new file mode 100644
index 0000000000..e5f1aa7b55
--- /dev/null
+++ b/eden/fs/py/eden/thrift/test/client_test.py
@@ -0,0 +1,26 @@
+#!/usr/bin/env python3
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This software may be used and distributed according to the terms of the
+# GNU General Public License version 2.
+
+import os.path
+import tempfile
+import unittest
+
+from eden.thrift.client import EdenNotRunningError, create_thrift_client
+
+
+class EdenClientTest(unittest.IsolatedAsyncioTestCase):
+    async def test_raise_EdenNotRunningError_when_no_socket(self):
+        with tempfile.TemporaryDirectory() as td:
+            sockname = os.path.join(td, "sock")
+            with self.assertRaises(EdenNotRunningError) as cm:
+                async with create_thrift_client(socket_path=sockname):
+                    pass
+            ex = cm.exception
+            self.assertEqual(sockname, ex.socket_path)
+            self.assertEqual(
+                "edenfs daemon does not appear to be running: tried " + sockname,
+                str(ex),
+            )