From afce59976bf864e2955d14118275f34e1d2e8306 Mon Sep 17 00:00:00 2001 From: Gary Wang Date: Tue, 29 Nov 2022 19:41:42 +0800 Subject: [PATCH 1/9] change: ignore VS Code settings files. --- .gitignore | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.gitignore b/.gitignore index 86cbf74..dec6755 100755 --- a/.gitignore +++ b/.gitignore @@ -32,5 +32,8 @@ MANIFEST .idea pydgraph.iml +# VS Code +.vscode + # Python Virtual Environments venv From a60774910cde4d44e8c746ee48d29034b19dd0fd Mon Sep 17 00:00:00 2001 From: Gary Wang Date: Tue, 29 Nov 2022 19:47:34 +0800 Subject: [PATCH 2/9] fix: absolute import for the test `helper` module. --- tests/test_acct_upsert.py | 2 +- tests/test_acl.py | 2 +- tests/test_async.py | 2 +- tests/test_client_stub.py | 2 +- tests/test_essentials.py | 2 +- tests/test_queries.py | 2 +- tests/test_txn.py | 2 +- tests/test_upsert_block.py | 2 +- tests/test_util.py | 2 +- 9 files changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/test_acct_upsert.py b/tests/test_acct_upsert.py index b79a2f6..f8cfe39 100644 --- a/tests/test_acct_upsert.py +++ b/tests/test_acct_upsert.py @@ -26,7 +26,7 @@ import pydgraph -from . import helper +from tests import helper CONCURRENCY = 5 FIRSTS = ['Paul', 'Eric', 'Jack', 'John', 'Martin'] diff --git a/tests/test_acl.py b/tests/test_acl.py index e62be6d..0bf26ba 100644 --- a/tests/test_acl.py +++ b/tests/test_acl.py @@ -23,7 +23,7 @@ import os import unittest -from . import helper +from tests import helper import pydgraph class TestACL(helper.ClientIntegrationTestCase): diff --git a/tests/test_async.py b/tests/test_async.py index f9f531f..6d14501 100644 --- a/tests/test_async.py +++ b/tests/test_async.py @@ -20,7 +20,7 @@ import json import pydgraph -from . import helper +from tests import helper class TestAsync(helper.ClientIntegrationTestCase): server_addr = 'localhost:9180' diff --git a/tests/test_client_stub.py b/tests/test_client_stub.py index 61e0a1a..a7dc33a 100644 --- a/tests/test_client_stub.py +++ b/tests/test_client_stub.py @@ -21,7 +21,7 @@ import sys import pydgraph -from . import helper +from tests import helper class TestDgraphClientStub(helper.ClientIntegrationTestCase): """Tests client stub.""" diff --git a/tests/test_essentials.py b/tests/test_essentials.py index 8b14453..1da4b16 100644 --- a/tests/test_essentials.py +++ b/tests/test_essentials.py @@ -21,7 +21,7 @@ import logging import json -from . import helper +from tests import helper class TestEssentials(helper.ClientIntegrationTestCase): diff --git a/tests/test_queries.py b/tests/test_queries.py index 2d50565..45a76e4 100755 --- a/tests/test_queries.py +++ b/tests/test_queries.py @@ -24,7 +24,7 @@ import pydgraph -from . import helper +from tests import helper class TestQueries(helper.ClientIntegrationTestCase): diff --git a/tests/test_txn.py b/tests/test_txn.py index d2a22e4..2b0974b 100644 --- a/tests/test_txn.py +++ b/tests/test_txn.py @@ -21,7 +21,7 @@ import pydgraph -from . import helper +from tests import helper class TestTxn(helper.ClientIntegrationTestCase): diff --git a/tests/test_upsert_block.py b/tests/test_upsert_block.py index 08f4472..39f5c45 100644 --- a/tests/test_upsert_block.py +++ b/tests/test_upsert_block.py @@ -20,7 +20,7 @@ import logging import json -from . import helper +from tests import helper class TestUpsertBlock(helper.ClientIntegrationTestCase): """Tests for Upsert Block""" diff --git a/tests/test_util.py b/tests/test_util.py index 1e3fd30..880d741 100644 --- a/tests/test_util.py +++ b/tests/test_util.py @@ -22,7 +22,7 @@ from pydgraph import util import pydgraph -from . import helper +from tests import helper class TestIsString(unittest.TestCase): From 828c1a298ea2b507569f2657fc2d1c09f2137c4a Mon Sep 17 00:00:00 2001 From: Gary Wang Date: Tue, 29 Nov 2022 19:52:54 +0800 Subject: [PATCH 3/9] change: allow `Txn` to accept commit kwargs. --- pydgraph/client.py | 4 ++-- pydgraph/txn.py | 10 ++++++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/pydgraph/client.py b/pydgraph/client.py index 0a3ef1c..bbb1091 100755 --- a/pydgraph/client.py +++ b/pydgraph/client.py @@ -150,9 +150,9 @@ def handle_alter_future(future): except Exception as error: DgraphClient._common_except_alter(error) - def txn(self, read_only=False, best_effort=False): + def txn(self, read_only=False, best_effort=False, **commit_kwargs): """Creates a transaction.""" - return txn.Txn(self, read_only=read_only, best_effort=best_effort) + return txn.Txn(self, read_only=read_only, best_effort=best_effort, **commit_kwargs) def any_client(self): """Returns a random gRPC client so that requests are distributed evenly among them.""" diff --git a/pydgraph/txn.py b/pydgraph/txn.py index 3c8f325..a828d93 100644 --- a/pydgraph/txn.py +++ b/pydgraph/txn.py @@ -42,7 +42,8 @@ class Txn(object): after calling commit. """ - def __init__(self, client, read_only=False, best_effort=False): + def __init__(self, client, read_only=False, best_effort=False, + timeout=None, metadata=None, credentials=None): if not read_only and best_effort: raise Exception('Best effort transactions are only compatible with ' 'read-only transactions') @@ -55,6 +56,11 @@ def __init__(self, client, read_only=False, best_effort=False): self._mutated = False self._read_only = read_only self._best_effort = best_effort + self._commit_kwargs = { + "timeout": timeout, + "metadata": metadata, + "credentials": credentials + } def query(self, query, variables=None, timeout=None, metadata=None, credentials=None, resp_format="JSON"): """Executes a query operation.""" @@ -164,7 +170,7 @@ def handle_mutate_future(txn, future, commit_now): response = future.result() except Exception as error: try: - txn.discard(timeout=timeout, metadata=metadata, credentials=credentials) + txn.discard(**txn._commit_kwargs) except: # Ignore error - user should see the original error. pass From 792df898905035a10264be291223d9e3633932af Mon Sep 17 00:00:00 2001 From: Gary Wang Date: Tue, 29 Nov 2022 19:56:25 +0800 Subject: [PATCH 4/9] fix: refer static method from `Txn` class name. --- pydgraph/txn.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pydgraph/txn.py b/pydgraph/txn.py index a828d93..662327c 100644 --- a/pydgraph/txn.py +++ b/pydgraph/txn.py @@ -158,7 +158,7 @@ def handle_query_future(future): try: response = future.result() except Exception as error: - txn._common_except_mutate(error) + Txn._common_except_mutate(error) return response @@ -174,7 +174,7 @@ def handle_mutate_future(txn, future, commit_now): except: # Ignore error - user should see the original error. pass - txn._common_except_mutate(error) + Txn._common_except_mutate(error) if commit_now: txn._finished = True From 735c5f81a9b9d21051c47e4ca4796f559b37b631 Mon Sep 17 00:00:00 2001 From: Gary Wang Date: Tue, 29 Nov 2022 20:07:40 +0800 Subject: [PATCH 5/9] enhancement: add context manager to `Txn` and `DgraphClientStub`. --- pydgraph/client.py | 21 +++++++++++++++++++++ pydgraph/client_stub.py | 33 +++++++++++++++++++++++++++++++++ pydgraph/txn.py | 12 ++++++++++++ tests/test_client_stub.py | 21 +++++++++++++++++++++ tests/test_txn.py | 30 ++++++++++++++++++++++++++++++ 5 files changed, 117 insertions(+) diff --git a/pydgraph/client.py b/pydgraph/client.py index bbb1091..86496b1 100755 --- a/pydgraph/client.py +++ b/pydgraph/client.py @@ -14,6 +14,7 @@ """Dgraph python client.""" +import contextlib import random from pydgraph import errors, txn, util @@ -164,3 +165,23 @@ def add_login_metadata(self, metadata): return new_metadata new_metadata.extend(metadata) return new_metadata + + @contextlib.contextmanager + def begin(self, + read_only:bool=False, best_effort:bool=False, + timeout = None, metadata = None, credentials = None): + '''Start a managed transaction. + + Note + ---- + Only use this function in ``with-as`` blocks. + ''' + tx = self.txn(read_only=read_only, best_effort=best_effort) + try: + yield tx + if read_only == False: + tx.commit(timeout=timeout, metadata=metadata, credentials=credentials) + except Exception as e: + raise e + finally: + tx.discard() \ No newline at end of file diff --git a/pydgraph/client_stub.py b/pydgraph/client_stub.py index 4034158..d4e7cf8 100644 --- a/pydgraph/client_stub.py +++ b/pydgraph/client_stub.py @@ -14,6 +14,7 @@ """Stub for RPC request.""" +import contextlib import grpc from pydgraph.meta import VERSION @@ -40,6 +41,14 @@ def __init__(self, addr='localhost:9080', credentials=None, options=None): self.channel = grpc.secure_channel(addr, credentials, options) self.stub = api_grpc.DgraphStub(self.channel) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() + if exc_type is not None: + raise exc_val def login(self, login_req, timeout=None, metadata=None, credentials=None): return self.stub.Login(login_req, timeout=timeout, metadata=metadata, @@ -111,3 +120,27 @@ def from_cloud(cloud_endpoint, api_key): client_stub = DgraphClientStub('{host}:{port}'.format( host=host, port="443"), composite_credentials, options=(('grpc.enable_http_proxy', 0),)) return client_stub + +@contextlib.contextmanager +def client_stub(addr='localhost:9080', **kwargs): + """ Create a managed DgraphClientStub instance. + + Parameters + ---------- + addr : str, optional + credentials : ChannelCredentials, optional + options: List[Dict] + An optional list of key-value pairs (``channel_arguments`` + in gRPC Core runtime) to configure the channel. + + Note + ---- + Only use this function in ``with-as`` blocks. + """ + stub = DgraphClientStub(addr=addr, **kwargs) + try: + yield stub + except Exception as e: + raise e + finally: + stub.close() \ No newline at end of file diff --git a/pydgraph/txn.py b/pydgraph/txn.py index 662327c..43a89c8 100644 --- a/pydgraph/txn.py +++ b/pydgraph/txn.py @@ -61,6 +61,18 @@ def __init__(self, client, read_only=False, best_effort=False, "metadata": metadata, "credentials": credentials } + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if exc_type is not None: + self.discard(**self._commit_kwargs) + raise exc_val + if self._read_only == False: + self.commit(**self._commit_kwargs) + else: + self.discard(**self._commit_kwargs) def query(self, query, variables=None, timeout=None, metadata=None, credentials=None, resp_format="JSON"): """Executes a query operation.""" diff --git a/tests/test_client_stub.py b/tests/test_client_stub.py index a7dc33a..8ab22b4 100644 --- a/tests/test_client_stub.py +++ b/tests/test_client_stub.py @@ -53,10 +53,31 @@ def test_close(self): client_stub.check_version(pydgraph.Check()) +class TestDgraphClientStubContextManager(helper.ClientIntegrationTestCase): + def setUp(self): + pass + + def test_context_manager(self): + with pydgraph.DgraphClientStub(addr=self.TEST_SERVER_ADDR) as client_stub: + ver = client_stub.check_version(pydgraph.Check()) + self.assertIsNotNone(ver) + + def test_context_manager_code_exception(self): + with self.assertRaises(AttributeError): + with pydgraph.DgraphClientStub(addr=self.TEST_SERVER_ADDR) as client_stub: + self.check_version(client_stub) # AttributeError: no such method + + def test_context_manager_function_wrapper(self): + with pydgraph.client_stub(addr=self.TEST_SERVER_ADDR) as client_stub: + ver = client_stub.check_version(pydgraph.Check()) + self.assertIsNotNone(ver) + + def suite(): """Returns a test suite object.""" suite_obj = unittest.TestSuite() suite_obj.addTest(TestDgraphClientStub()) + suite_obj.addTest(TestDgraphClientStubContextManager()) return suite_obj diff --git a/tests/test_txn.py b/tests/test_txn.py index 2b0974b..a15c5de 100644 --- a/tests/test_txn.py +++ b/tests/test_txn.py @@ -562,10 +562,40 @@ def test_sp_star2(self): self.assertEqual([{'uid': uid1}], json.loads(resp.json).get('me')) +class TestContextManager(helper.ClientIntegrationTestCase): + def setUp(self): + self.stub = pydgraph.DgraphClientStub(self.TEST_SERVER_ADDR) + self.client = pydgraph.DgraphClient(self.stub) + self.q = ''' + { + company(func: type(x.Company), first: 10){ + expand(_all_) + } + } + ''' + def tearDown(self) -> None: + self.stub.close() + + def test_context_manager_by_contextlib(self): + with self.client.begin(read_only=True, best_effort=True) as tx: + response = tx.query(self.q) + self.assertIsNotNone(response) + data = json.loads(response.json) + print(data) + + def test_context_manager_by_class(self): + with pydgraph.Txn(self.client, read_only=True, best_effort=True) as tx: + response = tx.query(self.q) + self.assertIsNotNone(response) + data = json.loads(response.json) + print(data) + + def suite(): s = unittest.TestSuite() s.addTest(TestTxn()) s.addTest(TestSPStar()) + s.addTest(TestContextManager()) return s From 522e0621c6f3ef3fa8490e5482b3d82110ef94e0 Mon Sep 17 00:00:00 2001 From: Gary Wang Date: Tue, 29 Nov 2022 20:54:52 +0800 Subject: [PATCH 6/9] docs: usage of Txn with context manager. --- README.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/README.md b/README.md index 4f04f52..c1726a3 100644 --- a/README.md +++ b/README.md @@ -251,6 +251,29 @@ finally: txn.discard() ``` +#### Using Transaction with Context Manager + +The Python context manager will automatically perform the "`commit`" action +after all queries and mutations have been done, and perform "`discard`" action +to clean the transaction. +When something goes wrong in the scope of context manager, "`commit`" will not +be called,and the "`discard`" action will be called to drop any potential changes. + +```python +with client.begin(read_only=False, best_effort=False) as txn: + # Do some queries or mutations here +``` + +or you can directly create a transaction from the `Txn` class. + +```python +with pydgraph.Txn(client, read_only=False, best_effort=False) as txn: + # Do some queries or mutations here +``` + +> `client.begin()` can only be used with "`with-as`" blocks, while `pydgraph.Txn` class can be directly called to instantiate a transaction object. + + ### Running a Query You can run a query by calling `Txn#query(string)`. You will need to pass in a From 1e6c7348457e93e2a5bc881b489be258d04d6341 Mon Sep 17 00:00:00 2001 From: Gary Wang Date: Tue, 29 Nov 2022 21:02:38 +0800 Subject: [PATCH 7/9] docs: use context manager to clean resources. --- README.md | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/README.md b/README.md index c1726a3..81a61bb 100644 --- a/README.md +++ b/README.md @@ -393,6 +393,24 @@ stub1.close() stub2.close() ``` +#### Use context manager to automatically clean resources + +Use function call: + +```python +with pydgraph.client_stub(SERVER_ADDR) as stub1: + with pydgraph.client_stub(SERVER_ADDR) as stub2: + client = pydgraph.DgraphClient(stub1, stub2) +``` + +Use class constructor: + +```python +with pydgraph.DgraphClientStub(SERVER_ADDR) as stub1: + with pydgraph.DgraphClientStub(SERVER_ADDR) as stub2: + client = pydgraph.DgraphClient(stub1, stub2) +``` + ### Setting Metadata Headers Metadata headers such as authentication tokens can be set through the metadata of gRPC methods. Below is an example of how to set a header named "auth-token". From 162fd0ee440838e2ff6daba155d5376c04e67013 Mon Sep 17 00:00:00 2001 From: Gary Wang Date: Tue, 29 Nov 2022 21:16:50 +0800 Subject: [PATCH 8/9] update: explain the `client`'s effective scope. --- README.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/README.md b/README.md index 81a61bb..491ae74 100644 --- a/README.md +++ b/README.md @@ -411,6 +411,10 @@ with pydgraph.DgraphClientStub(SERVER_ADDR) as stub1: client = pydgraph.DgraphClient(stub1, stub2) ``` +Note: `client` should be used inside the "`with-as`" block. The resources related to +`client` will be automatically released outside the block and `client` is not usable +any more. + ### Setting Metadata Headers Metadata headers such as authentication tokens can be set through the metadata of gRPC methods. Below is an example of how to set a header named "auth-token". From 20b4da734673e2dcdda647745b0b9e547b431613 Mon Sep 17 00:00:00 2001 From: Gary Wang Date: Wed, 30 Nov 2022 07:54:24 +0800 Subject: [PATCH 9/9] fix: avoid repeat commit when Txn finished. --- pydgraph/client.py | 2 +- pydgraph/txn.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pydgraph/client.py b/pydgraph/client.py index 86496b1..03aa4ac 100755 --- a/pydgraph/client.py +++ b/pydgraph/client.py @@ -179,7 +179,7 @@ def begin(self, tx = self.txn(read_only=read_only, best_effort=best_effort) try: yield tx - if read_only == False: + if read_only == False and tx._finished == False: tx.commit(timeout=timeout, metadata=metadata, credentials=credentials) except Exception as e: raise e diff --git a/pydgraph/txn.py b/pydgraph/txn.py index 43a89c8..20e5dea 100644 --- a/pydgraph/txn.py +++ b/pydgraph/txn.py @@ -69,7 +69,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): if exc_type is not None: self.discard(**self._commit_kwargs) raise exc_val - if self._read_only == False: + if self._read_only == False and self._finished == False: self.commit(**self._commit_kwargs) else: self.discard(**self._commit_kwargs)