RosCassandra.py

Go to the documentation of this file.

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import roslib; roslib.load_manifest("cassandra_ros")
import roslib.message
import rospy
import rosmsg
import rostopic
#import rosjson
import genpy.dynamic

import rosbag
import simplejson as json
import hashlib
import time
import StringIO

import cql
import pycassa

from Cassandra import Cassandra
from CassandraTopic import CassandraTopic

class RosCassandraException(rospy.ROSException): pass

class RosCassandra(Cassandra):
    
    def __init__(self, host="localhost", port=9160, keyspace="ros"):
        Cassandra.__init__(self, host, port)
        
        self.cql_conn   = None
        self.cql_cursor = None        
    
    def addTopic(self, topic, cassandra_format=None,     # binary, json-string, ros...
                               msg_class=None,            # tfMessage, uint8, etc
                               msg_package=None,          # the package where the msg-format was defined
                               key_format=None,           # Timestamp, hash of msg or a part of the msg, like a sequence 
                               key_msg_part=None,         # if a part of the message was defined ... wich one
                               comment = "",              # add a comment to topic
                               date=None):
        
        if self.existTopic(topic):
            return False
        
        topic_md5 = self.topic2Hash(topic)
        
        #column_validation_classes={'data':pycassa.BYTES_TYPE,
        #                           'format':pycassa.UTF8_TYPE,
        #                           'header.frame_id':pycassa.UTF8_TYPE,
        #                           'header.seq':pycassa.INT_TYPE,
        #                           'header.stamp.nsecs':pycassa.INT_TYPE,
        #                           'header.stamp.secs':pycassa.INT_TYPE}
        
        casTopic = CassandraTopic(topic, None, None,
                                  cassandra_format,
                                  msg_class, msg_package,
                                  key_format, key_msg_part,
                                  comment, date)
        
        column_validation_classes = casTopic.getColumnValidationClasses()
               
        column = self.createColumn(topic_md5, column_validation_classes=column_validation_classes)
        
        meta = casTopic.getMeta()
        self.setTopicMeta(topic, meta)
        
    def topic2Hash(self, topic):
        return hashlib.md5(topic).hexdigest()
        
    def getTopic(self, topic):
        
        if not self.existTopic(topic):
            return False        
        
        topic_md5 = self.topic2Hash(topic)
        column = self.getColumn(topic_md5)
        
        meta = self.getTopicMeta(topic)
        
        return CassandraTopic(column=column, cursor=None, meta=meta)
    
    def removeTopic(self, topic):
        if self.existTopic(topic):
            topic_md5 = self.topic2Hash(topic)
            self.dropColumn(topic_md5)
        
    def getTopicMeta(self, topic=None, topic_md5=None):
        if topic:
            topic_md5 = self.topic2Hash(topic)
        try:
            meta = json.loads(self.getColumnComment(topic_md5))
        except:
            meta = {}
        
        return meta
    
    def setTopicMeta(self, topic, meta):
        topic_md5 = self.topic2Hash(topic)
        self.setColumnComment(topic_md5, json.dumps(meta))
        
    def existTopic(self, topic):
        topic_md5 = self.topic2Hash(topic)
        return self.existColumn(topic_md5)
    
    def getAllTopics(self):
        topics = []
        columns = self.getAllColumns()

        for col in columns:
            meta = self.getTopicMeta(topic_md5=col)
            if type(meta) == dict:
                if meta.has_key('topic'):
                    topics.append(meta['topic'])
        
        return topics
     
    def renameTopic(self, topic, topic_new):
        topic_md5 = self.topic2Hash(topic)
        meta = self.getTopicMeta(topic_md5)
        meta['topic'] = topic_new
        
        pass
    
    
    def fileExport(self, topics, filename):
        pass
    
    def fileImport(self, filename, topics=None, format='ros', key='time', key_msg_part=None):
        print "read ...",
        bag = rosbag.Bag(filename)
        print "done"
        casTopic = {}
        i = 0
        for topic, msg, t in bag.read_messages(topics):
            if not casTopic.has_key(topic):
                if not self.existTopic(topic):
                    
                    self.addTopic(topic,
                                  cassandra_format=format,
                                  msg_class = msg.__class__._type.split('/')[1],
                                  msg_package = msg.__class__._type.split('/')[0],
                                  key_format=key, 
                                  key_msg_part=key_msg_part,
                                  comment = "",
                                  date=None)
                    
                casTopic[topic] = self.getTopic(topic)

            i += 1
            print i
            #time.sleep(5)
            if key == 'time':
                casTopic[topic].addData(msg, str(t.to_nsec()))
            else:
                casTopic[topic].addData(msg)
        
    
    def countTopicData(self, topic):
        try:
            topic_md5 = self.topic2Hash(topic)
            self.cursor.execute("select count(*) from '"+topic_md5+"'")
            return self.cursor.fetchone()[0]
        except:
            return 0
    
    def exequteCQL(self, query):
        # establish the cql connection
        if self.cql_cursor == None:
            self.cql_conn = cql.connect(self.host, self.port, self.keyspace)
            self.cql_cursor = self.cql_conn.cursor()
        
        topics = self.getAllTopics()
        
        # rename topics to column names ...
        for topic in topics:
            query = query.replace('"'+topic+'"', '"'+self.topic2Hash(topic)+'"')
        print query
        
        #return query
        # execute querie
        self.cql_cursor.execute(query)
        return self.cql_cursor.fetchall()
    
    def createIndex(self, topic, column):
        indexname = topic+column
        indexname = indexname.replace(".","_")
        column_validator = self.getTopic(topic).column.column_validators[column]
        
        self.sysManager.create_index(self.keyspace,
                                     self.topic2Hash(topic), 
                                     column, 
                                     column_validator,
                                     index_name=indexname)

    def removeIndex(self, topic, column):
        self.sysManager.drop_index(self.keyspace, self.topic2Hash(topic), column)