adds testCities

Author: WolfgangFahl

dg/dgraph.py

@@ -5,20 +5,22 @@
 '''
 import pydgraph
 import json
+import time
 
 class Dgraph(object):
     '''
     wrapper for https://dgraph.io/ GraphQL database
     '''
 
-    def __init__(self, host='localhost',port=9080, debug=False):
+    def __init__(self, host='localhost',port=9080, debug=False,profile=False):
         '''
         Constructor
         '''
         self.host=host
         self.port=port
 
         self.debug=debug
+        self.profile=profile
         self.client_stub = pydgraph.DgraphClientStub('%s:%d' % (host,port))
         self.client =pydgraph.DgraphClient(self.client_stub)
 
@@ -29,32 +31,57 @@ def addSchema(self,schema):
         result=self.client.alter(pydgraph.Operation(schema=schema))
         return result
 
-    def addData(self,obj=None,nquads=None):
+    def addData(self,obj=None,nquads=None,limit=None, batchSize=None):
         '''
-        add the given object or nquads if  both are given nquads are ignored
+        add the given object/list of objects or nquads if  both are given nquads are ignored
         '''
+        itemList=obj
+        if itemList is None:
+            self.addDataTxn(obj=None,nquads=nquads,itemTitle="nquads")
+        else:         
+            if type(itemList) is list:
+                if limit is not None:
+                    itemList=itemList[:limit]
+                if batchSize is None:
+                    self.addDataTxn(obj=itemList)
+                else:
+                    startTime=time.time()
+                    # store the list in batches
+                    for i in range(0, len(itemList), batchSize):
+                        itemBatch=itemList[i:i+batchSize]
+                        self.addDataTxn(obj=itemBatch, title="batch")
+                    print("addData for %9d items in %6.1f secs" % (len(itemList),time.time()-startTime))
+            
+        
+    def addDataTxn(self,obj=None,nquads=None,title="addData",itemTitle="items"):    
         response=None
         # Create a new transaction.
         txn = self.client.txn()
 
         try:
+            itemList=obj
+            size=1
+            startTime=time.time()
             # Run mutation.
-            if obj is not None:
+            if itemList is not None:
                 # check whether obj is  a list of items 
                 # if do a mutation for every item in the list
-                if type(obj) is list:
-                    for item in obj:
+                if type(itemList) is list:
+                    size=len(itemList)
+                    for item in itemList:
                         txn.mutate(set_obj=item)
                 else:        
                     # single object
                     response = txn.mutate(set_obj=obj)
+                if self.profile:    
+                    print("%7s for %9d %s in %6.1f secs" % (title,size,itemTitle,time.time()-startTime))    
             if nquads is not None:
                 response = txn.mutate(set_nquads=nquads)    
             # Commit transaction.
             txn.commit()
         finally:
             if self.debug:
-                print(obj)
+                print(itemList)
             # Clean up. Calling this after txn.commit() is a no-op and hence safe.
             txn.discard()
         return response

scripts/dgraph

@@ -62,12 +62,13 @@ error() {
 # show usage
 #
 usage() {
-  echo "$0 [-b|--bash|-h|--help|-k|--kill|-p|--pull]"
+  echo "$0 [-b|--bash|-c|--clean|-h|--help|-k|--kill|-p|--pull]"
   echo ""
   echo "-b | --bash: start a bash terminal shell within the currently running container"
   echo "-h | --help: show this usage"
   echo "-k | --kill: stop the docker image"
   echo "-p | --pull: pull the docker image"
+  echo "-c | --clean: clean start with kill and purge of all data"
   exit 1
 }
 
@@ -112,7 +113,7 @@ zero() {
 # dgraph alpha
 #
 alpha() {
-  $docker exec $it dgraph dgraph alpha --lru_mb 2048 --zero localhost:5080 --whitelist 0.0.0.0/0
+  $docker exec $it dgraph dgraph alpha --lru_mb 4096 --zero localhost:5080 --whitelist 0.0.0.0/0
 }
 
 #
@@ -181,6 +182,16 @@ do
     alpha) alpha;;
     ratel) ratel;;
     zero) zero;;
+    -c|--clean)
+       stopImage
+       if [ -d ~/dgraph ]
+       then
+          color_msg $blue "removing content from ~/dgraph"
+          rm -rf ~/dgraph/*
+       fi
+       # call me with no parameters to start terminals
+       $me
+       ;;
     -p|--pull) pullImage;;
     -b|--bash) bashInto;;
     -k|--kill) stopImage;;

tests/testDgraph.py

@@ -19,6 +19,58 @@ def setUp(self):
 
     def tearDown(self):
         pass
+    
+    def testCities(self):
+        '''
+        test a list of cities
+        '''
+        cityJsonUrl="https://raw.githubusercontent.com/lutangar/cities.json/master/cities.json"
+        with urllib.request.urlopen(cityJsonUrl) as url:
+            cityList=json.loads(url.read().decode())
+        self.assertEqual(128769,(len(cityList)))
+        cityIter=iter(cityList)
+        limit=2000
+        for i in range(limit):
+            city=next(cityIter)
+            city['dgraph.type']='City'
+            lat=city['lat']
+            lng=city['lng']
+            #city['location']={'type': 'Point', 'coordinates': [lng,lat] }
+            #print("%d: %s" % (i,city))
+        dgraph=Dgraph(profile=True)
+        dgraph.drop_all()
+        schema='''
+name: string @index(exact) .
+country: string .   
+lat: float .
+lng: float .
+location: geo .
+type City {
+   name
+   country
+   lat
+   lng
+   location
+}'''
+        dgraph.addSchema(schema)
+        dgraph.addData(obj=cityList,limit=limit,batchSize=250)
+        query='''{ 
+  # get cities
+  cities(func: has(name)) {
+        country
+        name
+        lat
+        lng
+        location
+  }
+}
+        '''
+        queryResult=dgraph.query(query)
+        self.assertTrue('cities' in queryResult)
+        qCityList=queryResult['cities']
+        self.assertEqual(limit,qCityList)
+        dgraph.close()
+            
 
     def testCountries(self):
         '''