From b724c6f0f7741e80c1c6359e81b4feb97fe082cf Mon Sep 17 00:00:00 2001
From: Joohwan Oh <joohwan.oh@outlook.com>
Date: Sat, 4 Feb 2017 21:41:06 -0800
Subject: [PATCH] Add Database.get_document method to allow direct retrieval of
 documents by ID

---
 arango/cursor.py       |  2 ++
 arango/database.py     | 34 ++++++++++++++++++++++++++++++++++
 tests/test_cursor.py   | 17 ++++++++++++++++-
 tests/test_document.py | 38 +++++++++++++++++++++++++++++++++++++-
 4 files changed, 89 insertions(+), 2 deletions(-)

diff --git a/arango/cursor.py b/arango/cursor.py
index ddcb4d46..9119833f 100644
--- a/arango/cursor.py
+++ b/arango/cursor.py
@@ -38,6 +38,8 @@ def __exit__(self, *_):
         self.close(ignore_missing=True)
 
     def __repr__(self):
+        if self.id is None:
+            return '<ArangoDB cursor>'
         return '<ArangoDB cursor {}>'.format(self.id)
 
     @property
diff --git a/arango/database.py b/arango/database.py
index f4a7c52b..b7bd5db2 100644
--- a/arango/database.py
+++ b/arango/database.py
@@ -156,6 +156,40 @@ def properties(self):
         result['system'] = result.pop('isSystem')
         return result
 
+    def get_document(self, id, rev=None, match_rev=True):
+        """Retrieve a document by its ID (collection/key)
+
+        :param id: the document ID
+        :type id: str | unicode
+        :returns: the document or ``None`` if the document is missing
+        :rtype: dict
+        :param rev: the revision to compare with that of the retrieved document
+        :type rev: str | unicode
+        :param match_rev: if ``True``, check if the given revision and
+            the target document's revisions are the same, otherwise check if
+            the revisions are different (this flag has an effect only when
+            **rev** is given)
+        :type match_rev: bool
+        :raises arango.exceptions.DocumentRevisionError: if the given revision
+            does not match the revision of the retrieved document
+        :raises arango.exceptions.DocumentGetError: if the document cannot
+            be retrieved from the collection
+        """
+        res = self._conn.get(
+            '/_api/document/{}'.format(id),
+            headers=(
+                {'If-Match' if match_rev else 'If-None-Match': rev}
+                if rev is not None else {}
+            )
+        )
+        if res.status_code in {304, 412}:
+            raise DocumentRevisionError(res)
+        elif res.status_code == 404 and res.error_code == 1202:
+            return None
+        elif res.status_code in HTTP_OK:
+            return res.body
+        raise DocumentGetError(res)
+
     #########################
     # Collection Management #
     #########################
diff --git a/tests/test_cursor.py b/tests/test_cursor.py
index fb493804..054b14c3 100644
--- a/tests/test_cursor.py
+++ b/tests/test_cursor.py
@@ -294,4 +294,19 @@ def test_cursor_context_manager():
         assert clean_keys(cursor.__next__()) == doc1
     with pytest.raises(CursorCloseError):
         cursor.close(ignore_missing=False)
-    assert cursor.close(ignore_missing=True) is False
\ No newline at end of file
+    assert cursor.close(ignore_missing=True) is False
+
+
+@pytest.mark.order12
+def test_cursor_repr_no_id():
+    col.truncate()
+    col.import_bulk([doc1, doc2, doc3, doc4])
+    cursor = db.aql.execute(
+        'FOR d IN {} RETURN d'.format(col_name),
+        count=True,
+        batch_size=2,
+        ttl=1000,
+        optimizer_rules=['+all']
+    )
+    getattr(cursor, '_data')['id'] = None
+    assert repr(cursor) == '<ArangoDB cursor>'
diff --git a/tests/test_document.py b/tests/test_document.py
index 569439b6..93aaaf3c 100644
--- a/tests/test_document.py
+++ b/tests/test_document.py
@@ -1033,7 +1033,8 @@ def test_get():
     with pytest.raises(ArangoError):
         col.get('5', rev='bad_rev')
 
-    # Test get with correct revision and match_rev turned off
+    # TODO uncomment once match_rev flag is fixed
+    # # Test get with correct revision and match_rev turned off
     # bad_rev = col['5']['_rev'] + '000'
     # result = col.get('5', rev=bad_rev, match_rev=False)
     # assert result['_key'] == '5'
@@ -1051,6 +1052,41 @@ def test_get():
         iter(bad_col)
 
 
+def test_get_from_db():
+    # Set up test documents
+    col.import_bulk(test_docs)
+
+    # Test get existing document
+    result = db.get_document(col_name + '/1')
+    assert result['_key'] == '1'
+    assert result['val'] == 100
+
+    # Test get another existing document
+    result = db.get_document(col_name + '/2')
+    assert result['_key'] == '2'
+    assert result['val'] == 100
+
+    # Test get missing document
+    assert db.get_document(col_name + '/6') is None
+
+    # Test get with correct revision
+    good_rev = db.get_document(col_name + '/5')['_rev']
+    result = db.get_document(col_name + '/5', rev=good_rev)
+    assert result['_key'] == '5'
+    assert result['val'] == 300
+
+    # Test get with invalid revision
+    bad_rev = db.get_document(col_name + '/5')['_rev'] + '000'
+    with pytest.raises(ArangoError):
+        db.get_document(col_name + '/5', rev=bad_rev, match_rev=True)
+    with pytest.raises(ArangoError):
+        db.get_document(col_name + '/5', rev="bad_rev")
+
+    # Test get with missing collection
+    with pytest.raises(DocumentGetError):
+        _ = db.get_document(bad_col_name + '/1')
+
+
 def test_get_many():
     # Test precondition
     assert len(col) == 0