Add client wrapper for import of documents.

Author: kishorenc

examples/collection_operations.py

@@ -1,3 +1,4 @@
+import json
 import os
 import sys
 
@@ -7,31 +8,30 @@
 import typesense
 
 client = typesense.Client({
-  'api_key': 'abcd',
-  'nodes': [{
-    'host': 'localhost',
-    'port': '8108',
-    'protocol': 'http'
-  }],
-  'timeout_seconds': 2
+    'api_key': 'abcd',
+    'nodes': [{
+        'host': 'localhost',
+        'port': '8108',
+        'protocol': 'http'
+    }],
+    'timeout_seconds': 2
 })
 
-
 # Create a collection
 
 create_response = client.collections.create({
-  "name": "books",
-  "fields": [
-    {"name": "title", "type": "string" },
-    {"name": "authors", "type": "string[]" },
-    {"name": "authors_facet", "type": "string[]", "facet": True },
-    {"name": "publication_year", "type": "int32" },
-    {"name": "publication_year_facet", "type": "string", "facet": True },
-    {"name": "ratings_count", "type": "int32" },
-    {"name": "average_rating", "type": "float" },
-    {"name": "image_url", "type": "string" }
-  ],
-  "default_sorting_field": "ratings_count"
+    "name": "books",
+    "fields": [
+        {"name": "title", "type": "string"},
+        {"name": "authors", "type": "string[]"},
+        {"name": "authors_facet", "type": "string[]", "facet": True},
+        {"name": "publication_year", "type": "int32"},
+        {"name": "publication_year_facet", "type": "string", "facet": True},
+        {"name": "ratings_count", "type": "int32"},
+        {"name": "average_rating", "type": "float"},
+        {"name": "image_url", "type": "string"}
+    ],
+    "default_sorting_field": "ratings_count"
 })
 
 print(create_response)
@@ -48,18 +48,19 @@
 # Add a book
 
 hunger_games_book = {
-  'id': '1', 'original_publication_year': 2008, 'authors': ['Suzanne Collins'], 'average_rating': 4.34,
-  'publication_year': 2008, 'publication_year_facet': '2008', 'authors_facet': ['Suzanne Collins'],
-  'title': 'The Hunger Games',
-  'image_url': 'https://images.gr-assets.com/books/1447303603m/2767052.jpg', 
-  'ratings_count': 4780653 
+    'id': '1', 'original_publication_year': 2008, 'authors': ['Suzanne Collins'], 'average_rating': 4.34,
+    'publication_year': 2008, 'publication_year_facet': '2008', 'authors_facet': ['Suzanne Collins'],
+    'title': 'The Hunger Games',
+    'image_url': 'https://images.gr-assets.com/books/1447303603m/2767052.jpg',
+    'ratings_count': 4780653
 }
 
 client.collections['books'].documents.create(hunger_games_book)
 
 # Export the documents from a collection
 
-print(client.collections['books'].documents.export())
+exported_doc_strs = client.collections['books'].documents.export()
+print(exported_doc_strs)
 
 # Fetch a document in a collection
 
@@ -77,6 +78,14 @@
 
 print(client.collections['books'].documents['1'].delete())
 
+# Import documents into a collection
+docs_to_import = []
+for exported_doc_str in exported_doc_strs:
+    docs_to_import.append(json.loads(exported_doc_str))
+
+import_res = client.collections['books'].documents.create_many(docs_to_import)
+print(import_res["success"])
+
 # Drop the collection
 
 drop_response = client.collections['books'].delete()

examples/indexing_in_loop.py renamed to examples/index_and_search.py

@@ -48,7 +48,19 @@
 
 print(create_response)
 
+# Let's use bulk import
+book_documents = []
 with open('/tmp/books.jsonl') as infile:
     for json_line in infile:
-        book_document = json.loads(json_line)
-        client.collections['books'].documents.create(book_document)
+        book_documents.append(json.loads(json_line))
+
+print(client.collections['books'].documents.create_many(book_documents))
+
+i = 0
+while i < 10000:
+    res = client.collections['books'].documents.search({
+        'q': 'the',
+        'query_by': 'title',
+        'sort_by': 'ratings_count:desc'
+    })
+    print(res['found'])

typesense/api_call.py

@@ -1,3 +1,4 @@
+import json
 import time
 
 import requests
@@ -69,15 +70,19 @@ def make_request(self, fn, method, endpoint, as_json, **kwargs):
 
             try:
                 url = node.url() + endpoint
+                if kwargs.get('data') and not isinstance(kwargs['data'], str):
+                    kwargs['data'] = json.dumps(kwargs['data'])
+
                 r = fn(url, headers={ApiCall.API_KEY_HEADER_NAME: self.config.api_key}, **kwargs)
 
                 if 0 < r.status_code < 500:
                     node.healthy = True
 
-                if (method != 'post' and r.status_code != 200) or (method == 'post' and r.status_code != 201):
+                if (method != 'post' and r.status_code != 200) or \
+                   (method == 'post' and not (r.status_code == 200 or r.status_code == 201)):
                     error_message = r.json().get('message', 'API error.')
                     # print('error_message: ' + error_message)
-                    raise ApiCall.get_exception(r.status_code)(error_message)
+                    raise ApiCall.get_exception(r.status_code)(r.status_code, error_message)
 
                 return r.json() if as_json else r.text
             except requests.exceptions.Timeout:
@@ -102,11 +107,11 @@ def get(self, endpoint, params=None, as_json=True):
 
     def post(self, endpoint, body):
         return self.make_request(requests.post, 'post', endpoint, True,
-                                 json=body, timeout=self.config.timeout_seconds)
+                                 data=body, timeout=self.config.timeout_seconds)
 
     def put(self, endpoint, body):
         return self.make_request(requests.put, 'put', endpoint, True,
-                                 json=body, timeout=self.config.timeout_seconds)
+                                 data=body, timeout=self.config.timeout_seconds)
 
     def delete(self, endpoint):
         return self.make_request(requests.delete, 'delete', endpoint, True,

typesense/documents.py

@@ -1,3 +1,5 @@
+import json
+
 from .document import Document
 from .api_call import ApiCall
 
@@ -27,6 +29,14 @@ def _endpoint_path(self, action=None):
     def create(self, document):
         return self.api_call.post(self._endpoint_path(), document)
 
+    def create_many(self, documents):
+        document_strs = []
+        for document in documents:
+            document_strs.append(json.dumps(document))
+
+        docs_import = '\n'.join(document_strs)
+        return self.api_call.post(self._endpoint_path('import'), docs_import)
+
     def export(self):
         api_response = self.api_call.get(self._endpoint_path('export'), {}, as_json=False)
         return api_response.split('\n')