Bulbs 0.3 gremlin.py

gremlin.py

# -*- coding: utf-8 -*-
#
# Copyright 2011 James Thornton (http://jamesthornton.com)
# BSD License (see LICENSE for details)
#
"""
An interface for executing Gremlin scripts on the resource.

"""
import os
import yaml
from string import Template
from utils import initialize_element

    
class Gremlin(object):
    """An interface for executing Gremlin scripts on the resource."""

    def __init__(self,resource):
        self.resource = resource
        # Registering it here isn't a good idea b/c you will
        # have a chicken and egg scneario if you try to use
        # the sripts inside Resource -- register it at the resource level
        #self.register_scripts("gremlin","gremlin.yaml")

    def query(self,script,**kwds):
        """
        Returns initialized results of an arbitrary Gremlin scripts
        run on the resource.

        :param script: Gremlin script to send to the resource.
        :param kwds: Resource-specific keyword params.

        """
        resp = self.resource.gremlin(script,**kwds)
        for result in resp.results:
            yield initialize_element(self.resource,result)
 
    def execute(self,script,**kwds):
        """
        Returns raw results of an arbitrary Gremlin script.

        :param script: Gremlin script to send to the resource. 
        :param kwds: Resource-specific keyword params.

        """
        resp = self.resource.gremlin(script,**kwds)
        return list(resp.results)

    def register_scripts(self,name,file_name):
        """Only load/register the scripts library once to reduce overhead."""
        scripts = self.resource.config.get_scripts(name)
        if not scripts:
            scripts = Scripts(file_name)
            self.resource.config.register_scripts(name,scripts)
    
    def get_scripts(self,name="gremlin"):
        return self.resource.registry.get_scripts(name)


class Scripts(object):
    """Load Gremlin scripts from a YAML source file."""

    def __init__(self,file_name="gremlin.yaml"):
        self.file_name = self._get_file_name(file_name)
        self.templates = self._load_templates()

    def get(self,name,params={}):
        """Return a Gremlin script, generated from the params."""
        template = self.templates.get(name)
        params = self._quote_params(params)
        return template.substitute(params)
        
    def refresh(self):
        """Refresh the stored templates from the YAML source."""
        self.templates = self._load_templates()

    def _load_templates(self):
        return self._parse_yaml(self.file_name)

    def _get_file_name(self,file_name):
        if file_name == "gremlin.yaml":
            dir_name = os.path.dirname(__file__)
            file_name = "%s/%s" % (dir_name,file_name)
        return file_name

    def _parse_yaml(self,file_name):
        templates = dict()
        f = open(file_name)
        yaml_map = yaml.load(f)    
        for name, template in yaml_map.items():
            #template = ';'.join(lines.split('\n'))
            templates[name] = Template(template)
        return templates

    def _quote_params(self,params):
        quoted_tuple = map(self._quote,params.items())
        params = dict(quoted_tuple)
        return params

    def _quote(self,pair):
        key, value = pair
        if type(value) == str:
            value = "'%s'" % value
        elif value is None:
            value = ""
        return key, value


class ScriptWriter(object):
    """
    ScriptWriter is an experiment that would be akin to a Python-based 
    Google Web Toolkit (http://code.google.com/webtoolkit/) for 
    building Gremlin scripts. And while it works well, I think it's simpler
    to source Gremlin code from gremlin.yaml. But I'm leaving ScriptWriter
    in the codebase for now so others can experiment with it. - James

    Example:

    def create_indexed_vertex(index_name,data,keys=None):
        s = ScriptWriter()
        s.start_transaction()
        s.add_vertex("v")
        s.set_property_data("v",data)
        s.get_index("i",index_name,"Vertex")
        keys = s.get_keys(data,keys)
        for key, value in data.items():
            if key in keys:
                s.index_put("i",key,value,"v")
        s.end_transaction()
        s.return_var("v")
        return s.get()

    data = dict(name="James",age=34)
    script = create_indexed_vertex("people",data)
    resp = self.resource.gremlin(script)
    """
    
    def __init__(self):
        self.lines = []

    def __add__(self,line):
        self.lines.append(line)

    # Elements
    def add_vertex(self,varname):
        line = "Vertex %s = g.addVertex(null)" % (varname)
        self.lines.append(line)

    def add_edge(self,varname,outV,label,inV):
        line  = "Edge %s = g.addEdge(null,%s,%s,%s)" % (varname,outV,inV,label)
        self.lines.append(line)

    def set_property(self,element,key,value):
        line = "%s.setProperty('%s',%s)" % (element,key,self.quote(value))
        self.lines.append(line)

    def set_property_data(self,element,data):
        for key, value in data.items():
            self.set_property(element,key,value)

    # Indices
    def get_index(self,varname,index_name,index_class):
        line = "%s = g.getIndex(%s,%s)" % (varname, index_name,index_class)
        self.lines.append(line)

    def index_put(self,index,key,value,element):
        line = "%s.put('%s',%s,%s)" % (index,key,self.quote(value),element)
        self.lines.append(line)

    def index_get(self,index,key,value):
        line = "%s.get('%s',%s,%s)" % (index,key,self.quote(value))
        
    def index_remove(self,index,key,value,element):
        line = "%s.remove('%s',%s,%s)" % (index,key,self.quote(value),element)
        self.lines.append(line)

    # Graph
    def return_var(self,varname):
        self.lines.append("return %s" % varname)

    def start_transaction(self,buffer_size=0):
        self.lines.append("g.setMaxBufferSize(%d)" % buffer_size)
        self.lines.append("g.startTransaction()")
        
    def end_transaction(self):
        self.lines.append("g.stopTransaction(TransactionalGraph.Conclusion.SUCCESS)")

    # Script Methods
    def append(self,script):
        self.lines.append(script)

    def get(self):
        script = ";".join(self.lines)
        return script

    def display(self):
        script = ";\n".join(self.lines)
        print script

    def compile(self):
        # store pre-compiled tempaltes in config?
        pass

    # Utils
    def quote(value):
        # quote it if it's a string, set to null if None, else return the value
        if type(value) == str:
            value = "'%s'" % value
        elif value is None:
            value = "null"
        return value

    def get_keys(data,keys):
        if not keys:
            keys = data.keys()
        return keys