add ability to create an index after the import (#38)

* add ability to create an index after the import

* fix typo and reword to be more clear

Author: maguec

README.md

@@ -54,6 +54,7 @@ python3 redisgraph_bulk_loader/bulk_insert.py GRAPHNAME [OPTIONS]
 |  -t   | --max-token-count INT      |            (Debug argument) Max number of tokens sent in each Redis query (default 1024)             |
 |  -b   | --max-buffer-size INT      |               (Debug argument) Max batch size (MBs) of each Redis query (default 4096)               |
 |  -c   | --max-token-size INT       |              (Debug argument) Max size (MBs) of each token sent to Redis (default 500)               |
+|  -i   | --index Label:Property     |              After bulk import, create an Index on provided Label:Property pair (optional)           |
 
 
 The only required arguments are the name to give the newly-created graph (which can appear anywhere) and at least one node CSV file.

redisgraph_bulk_loader/bulk_insert.py

@@ -62,7 +62,8 @@ def process_entities(entities):
 @click.option('--max-token-count', '-c', default=1024, help='max number of processed CSVs to send per query (default 1024)')
 @click.option('--max-buffer-size', '-b', default=2048, help='max buffer size in megabytes (default 2048)')
 @click.option('--max-token-size', '-t', default=500, help='max size of each token in megabytes (default 500, max 512)')
-def bulk_insert(graph, host, port, password, nodes, nodes_with_label, relations, relations_with_type, separator, enforce_schema, skip_invalid_nodes, skip_invalid_edges, quote, max_token_count, max_buffer_size, max_token_size):
+@click.option('--index', '-i', multiple=True, help='Label:Propery on which to create an index')
+def bulk_insert(graph, host, port, password, nodes, nodes_with_label, relations, relations_with_type, separator, enforce_schema, skip_invalid_nodes, skip_invalid_edges, quote, max_token_count, max_buffer_size, max_token_size, index):
     if sys.version_info[0] < 3:
         raise Exception("Python 3 is required for the RedisGraph bulk loader.")
 
@@ -115,6 +116,16 @@ def bulk_insert(graph, host, port, password, nodes, nodes_with_label, relations,
     end_time = timer()
     query_buf.report_completion(end_time - start_time)
 
+    for i in index:
+        l, p = i.split(":")
+        print("Creating Index on Label: %s, Property: %s" %(l, p))
+        try:
+            index_create = client.execute_command("GRAPH.QUERY", graph, "CREATE INDEX ON :%s(%s)" %(l, p))
+            for z in index_create:
+                print(z[0].decode("utf-8") )
+        except redis.exceptions.ResponseError as e:
+            print("Unable to create Index on Label: %s, Property %s" %(l, p))
+            print(e)
 
 if __name__ == '__main__':
     bulk_insert()

test/test_bulk_loader.py

@@ -660,6 +660,31 @@ def test16_error_on_schema_failure(self):
             self.assertEqual(sys.exc_info()[0].__name__, 'SchemaError')
             self.assertIn("Could not parse 'strval' as an array", e.args)
 
+    def test17_ensure_index_is_created(self):
+        graphname = "index_test"
+        with open('/tmp/nodes_index.tmp', mode='w') as csv_file:
+            out = csv.writer(csv_file, delimiter='|')
+            out.writerow(['name:STRING', 'age:INT'])
+            out.writerow(['Alex', 17])
+            out.writerow(['Sean', 12])
+        csv_file.close()
+
+        runner = CliRunner()
+        res = runner.invoke(bulk_insert, ['--nodes-with-label', 'Person', '/tmp/nodes_index.tmp',
+                                          '--separator', '|',
+                                          '--index', 'Person:age',
+                                          '--enforce-schema',
+                                          graphname], catch_exceptions=False)
+
+        self.assertEqual(res.exit_code, 0)
+        self.assertIn('2 nodes created', res.output)
+        self.assertIn('Indices created: 1', res.output)
+
+        graph = Graph(graphname, self.redis_con)
+        r = redis.Redis(host='localhost', port=6379, decode_responses=True)
+        res = r.execute_command("GRAPH.EXPLAIN", graphname, 'MATCH (p:Person) WHERE p.age > 16 RETURN p')
+        self.assertIn('        Index Scan | (p:Person)', res)
+
 
 if __name__ == '__main__':
     unittest.main()