tempest_db 0.13.2 → 0.14.0

data/lib/gen/tempest_types.rb

@@ -6,88 +6,204 @@
 
 require 'thrift'
 
-class InvalidNodeIdException < ::Thrift::Exception
-  include ::Thrift::Struct, ::Thrift::Struct_Union
-  def initialize(message=nil)
-    super()
-    self.message = message
-  end
+module Teapot
+  module TempestDB
+    module DegreeFilterTypes
+      INDEGREE_MIN = 0
+      INDEGREE_MAX = 1
+      OUTDEGREE_MIN = 2
+      OUTDEGREE_MAX = 3
+      VALUE_MAP = {0 => "INDEGREE_MIN", 1 => "INDEGREE_MAX", 2 => "OUTDEGREE_MIN", 3 => "OUTDEGREE_MAX"}
+      VALID_VALUES = Set.new([INDEGREE_MIN, INDEGREE_MAX, OUTDEGREE_MIN, OUTDEGREE_MAX]).freeze
+    end
+
+    class Node
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+      TYPE = 1
+      ID = 2
+
+      FIELDS = {
+        TYPE => {:type => ::Thrift::Types::STRING, :name => 'type'},
+        ID => {:type => ::Thrift::Types::STRING, :name => 'id'}
+      }
+
+      def struct_fields; FIELDS; end
+
+      def validate
+        raise ::Thrift::ProtocolException.new(::Thrift::ProtocolException::UNKNOWN, 'Required field type is unset!') unless @type
+        raise ::Thrift::ProtocolException.new(::Thrift::ProtocolException::UNKNOWN, 'Required field id is unset!') unless @id
+      end
+
+      ::Thrift::Struct.generate_accessors self
+    end
+
+    class MonteCarloPageRankParams
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+      NUMSTEPS = 1
+      RESETPROBABILITY = 2
+      MAXINTERMEDIATENODEDEGREE = 3
+      MINREPORTEDVISITS = 4
+      MAXRESULTCOUNT = 5
+
+      FIELDS = {
+        NUMSTEPS => {:type => ::Thrift::Types::I32, :name => 'numSteps'},
+        RESETPROBABILITY => {:type => ::Thrift::Types::DOUBLE, :name => 'resetProbability'},
+        MAXINTERMEDIATENODEDEGREE => {:type => ::Thrift::Types::I32, :name => 'maxIntermediateNodeDegree', :default => 1000, :optional => true},
+        MINREPORTEDVISITS => {:type => ::Thrift::Types::I32, :name => 'minReportedVisits', :optional => true},
+        MAXRESULTCOUNT => {:type => ::Thrift::Types::I32, :name => 'maxResultCount', :optional => true}
+      }
+
+      def struct_fields; FIELDS; end
+
+      def validate
+        raise ::Thrift::ProtocolException.new(::Thrift::ProtocolException::UNKNOWN, 'Required field numSteps is unset!') unless @numSteps
+        raise ::Thrift::ProtocolException.new(::Thrift::ProtocolException::UNKNOWN, 'Required field resetProbability is unset!') unless @resetProbability
+      end
+
+      ::Thrift::Struct.generate_accessors self
+    end
+
+    class BidirectionalPPRParams
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+      RELATIVEERROR = 1
+      RESETPROBABILITY = 2
+      MINPROBABILITY = 3
 
-  MESSAGE = 1
+      FIELDS = {
+        RELATIVEERROR => {:type => ::Thrift::Types::DOUBLE, :name => 'relativeError'},
+        RESETPROBABILITY => {:type => ::Thrift::Types::DOUBLE, :name => 'resetProbability'},
+        MINPROBABILITY => {:type => ::Thrift::Types::DOUBLE, :name => 'minProbability', :optional => true}
+      }
+
+      def struct_fields; FIELDS; end
+
+      def validate
+        raise ::Thrift::ProtocolException.new(::Thrift::ProtocolException::UNKNOWN, 'Required field relativeError is unset!') unless @relativeError
+        raise ::Thrift::ProtocolException.new(::Thrift::ProtocolException::UNKNOWN, 'Required field resetProbability is unset!') unless @resetProbability
+      end
+
+      ::Thrift::Struct.generate_accessors self
+    end
 
-  FIELDS = {
-    MESSAGE => {:type => ::Thrift::Types::STRING, :name => 'message'}
-  }
+    class InvalidNodeIdException < ::Thrift::Exception
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+      def initialize(message=nil)
+        super()
+        self.message = message
+      end
 
-  def struct_fields; FIELDS; end
+      MESSAGE = 1
 
-  def validate
-  end
+      FIELDS = {
+        MESSAGE => {:type => ::Thrift::Types::STRING, :name => 'message'}
+      }
 
-  ::Thrift::Struct.generate_accessors self
-end
+      def struct_fields; FIELDS; end
 
-class InvalidIndexException < ::Thrift::Exception
-  include ::Thrift::Struct, ::Thrift::Struct_Union
-  def initialize(message=nil)
-    super()
-    self.message = message
-  end
+      def validate
+      end
 
-  MESSAGE = 1
+      ::Thrift::Struct.generate_accessors self
+    end
 
-  FIELDS = {
-    MESSAGE => {:type => ::Thrift::Types::STRING, :name => 'message'}
-  }
+    class InvalidIndexException < ::Thrift::Exception
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+      def initialize(message=nil)
+        super()
+        self.message = message
+      end
 
-  def struct_fields; FIELDS; end
+      MESSAGE = 1
 
-  def validate
-  end
+      FIELDS = {
+        MESSAGE => {:type => ::Thrift::Types::STRING, :name => 'message'}
+      }
 
-  ::Thrift::Struct.generate_accessors self
-end
+      def struct_fields; FIELDS; end
 
-class InvalidArgumentException < ::Thrift::Exception
-  include ::Thrift::Struct, ::Thrift::Struct_Union
-  def initialize(message=nil)
-    super()
-    self.message = message
-  end
+      def validate
+      end
 
-  MESSAGE = 1
+      ::Thrift::Struct.generate_accessors self
+    end
 
-  FIELDS = {
-    MESSAGE => {:type => ::Thrift::Types::STRING, :name => 'message'}
-  }
+    class InvalidArgumentException < ::Thrift::Exception
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+      def initialize(message=nil)
+        super()
+        self.message = message
+      end
 
-  def struct_fields; FIELDS; end
+      MESSAGE = 1
 
-  def validate
-  end
+      FIELDS = {
+        MESSAGE => {:type => ::Thrift::Types::STRING, :name => 'message'}
+      }
 
-  ::Thrift::Struct.generate_accessors self
-end
+      def struct_fields; FIELDS; end
 
-class BidirectionalPPRParams
-  include ::Thrift::Struct, ::Thrift::Struct_Union
-  RELATIVEERROR = 1
-  RESETPROBABILITY = 2
-  MINPROBABILITY = 3
+      def validate
+      end
 
-  FIELDS = {
-    RELATIVEERROR => {:type => ::Thrift::Types::DOUBLE, :name => 'relativeError'},
-    RESETPROBABILITY => {:type => ::Thrift::Types::DOUBLE, :name => 'resetProbability'},
-    MINPROBABILITY => {:type => ::Thrift::Types::DOUBLE, :name => 'minProbability', :optional => true}
-  }
+      ::Thrift::Struct.generate_accessors self
+    end
 
-  def struct_fields; FIELDS; end
+    class UndefinedGraphException < ::Thrift::Exception
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+      def initialize(message=nil)
+        super()
+        self.message = message
+      end
 
-  def validate
-    raise ::Thrift::ProtocolException.new(::Thrift::ProtocolException::UNKNOWN, 'Required field relativeError is unset!') unless @relativeError
-    raise ::Thrift::ProtocolException.new(::Thrift::ProtocolException::UNKNOWN, 'Required field resetProbability is unset!') unless @resetProbability
-  end
+      MESSAGE = 1
 
-  ::Thrift::Struct.generate_accessors self
-end
+      FIELDS = {
+        MESSAGE => {:type => ::Thrift::Types::STRING, :name => 'message'}
+      }
+
+      def struct_fields; FIELDS; end
+
+      def validate
+      end
+
+      ::Thrift::Struct.generate_accessors self
+    end
 
+    class SQLException < ::Thrift::Exception
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+      def initialize(message=nil)
+        super()
+        self.message = message
+      end
+
+      MESSAGE = 1
+
+      FIELDS = {
+        MESSAGE => {:type => ::Thrift::Types::STRING, :name => 'message'}
+      }
+
+      def struct_fields; FIELDS; end
+
+      def validate
+      end
+
+      ::Thrift::Struct.generate_accessors self
+    end
+
+    class UnequalListSizeException < ::Thrift::Exception
+      include ::Thrift::Struct, ::Thrift::Struct_Union
+
+      FIELDS = {
+
+      }
+
+      def struct_fields; FIELDS; end
+
+      def validate
+      end
+
+      ::Thrift::Struct.generate_accessors self
+    end
+
+  end
+end

data/lib/tempest_db.rb

@@ -18,8 +18,8 @@
 $:.push(File.expand_path(File.join(File.dirname(__FILE__),"gen")))
 
 require 'tempest_d_b_service'
-require 'tempest_db_constants'
-require 'tempest_db_types'
+require 'tempest_constants'
+require 'tempest_types'
 require 'base_thrift_client'
 
 # May also require the types and the constants for more complex services
@@ -35,10 +35,14 @@ module Teapot
     class TempestClient
       @thrift_client = nil
 
+      # @param [String] server
+      # @param [Integer] port
       def initialize(server, port)
         @thrift_client = Teapot::TempestDB::TempestThriftClient.new(server, port)
       end
 
+      # @param [String] json_attribute
+      # @return [Object]
       def TempestClient.jsonToValue(json_attribute)
         if json_attribute[0] == '"'
           json_attribute[1...-1]
@@ -54,30 +58,53 @@ module Teapot
         end
       end
 
-      def get_multi_node_attribute(node_type, nodeIds, attribute_name)
-        id_to_json = @thrift_client.with_retries { |executor|
-          executor.getMultiNodeAttributeAsJSON(node_type, nodeIds, attribute_name)
+      # @param [Array<Teapot::TempestDB::Node>] nodes
+      # @param [String] attribute_name
+      # @return [Hash]
+      def get_multi_node_attribute(nodes, attribute_name)
+        node_to_json = @thrift_client.with_retries { |executor|
+          executor.getMultiNodeAttributeAsJSON(nodes, attribute_name)
         }
-        Hash[id_to_json.map{ |k,v| [k, TempestClient.jsonToValue(v)] }]
+        Hash[node_to_json.map{ |k,v| [k, TempestClient.jsonToValue(v)] }]
       end
 
-      def get_node_attribute(node_type, nodeId, attribute_name)
-        self.get_multi_node_attribute(node_type, [nodeId], attribute_name)[nodeId]
+      # @param [Teapot::TempestDB::Node] node
+      # @param [String] attribute_name
+      # @return [String]
+      def get_node_attribute(node, attribute_name)
+        self.get_multi_node_attribute([node], attribute_name)[node]
       end
 
-      def k_step_out_neighbors_filtered(edge_type, source_id, k, sql_clause: "", degree_filter: {}, alternating: true)
+      # @param [String] edge_type
+      # @param [Teapot::TempestDB::Node] source_node
+      # @param [Integer] k
+      # @param [String] sql_clause
+      # @param [Hash] degree_filter
+      # @param [Boolean] alternating
+      # @return [Teapot::TempestDB::Node]
+      def k_step_out_neighbors_filtered(edge_type, source_node, k, sql_clause: "", degree_filter: {}, alternating: true)
         @thrift_client.with_retries { |executor|
-          executor.kStepOutNeighborsFiltered(edge_type, source_id, k, sql_clause, degree_filter, alternating)
+          executor.kStepOutNeighborsFiltered(edge_type, source_node, k, sql_clause, degree_filter, alternating)
         }
       end
 
-      def k_step_in_neighbors_filtered(edge_type, source_id, k, sql_clause: "", degree_filter: {}, alternating: true)
+      # @param [String] edge_type
+      # @param [Teapot::TempestDB::Node] source_node
+      # @param [Integer] k
+      # @param [String] sql_clause
+      # @param [Hash] degree_filter
+      # @param [Boolean] alternating
+      # @return [Teapot::TempestDB::Node]
+      def k_step_in_neighbors_filtered(edge_type, source_node, k, sql_clause: "", degree_filter: {}, alternating: true)
         @thrift_client.with_retries { |executor|
-          executor.kStepInNeighborsFiltered(edge_type, source_id, k, sql_clause, degree_filter, alternating)
+          executor.kStepInNeighborsFiltered(edge_type, source_node, k, sql_clause, degree_filter, alternating)
         }
       end
 
-      # Convenience method to return the unique node id satisfying a condition
+      # Convenience method to return the unique node node satisfying a condition
+      # @param [String] node_type
+      # @param [String] sql_clause
+      # @return [Teapot::TempestDB::Node]
       def unique_node(node_type, sql_clause)
         matching_nodes = @thrift_client.with_retries { |executor|
           executor.nodes(node_type, sql_clause)
@@ -91,91 +118,164 @@ module Teapot
 
       # TODO: For methods that only need to be included (for irb auto-complete) and converted from camelCase to
       # snake_case, is there a simple way of automating this?
-      def out_degree(edge_type, id)
+      # @param [String] edge_type
+      # @param [Teapot::TempestDB::Node] node
+      # @return [Integer]
+      def out_degree(edge_type, node)
         @thrift_client.with_retries { |executor|
-          executor.outDegree(edge_type, id)
+          executor.outDegree(edge_type, node)
         }
       end
 
-      def in_degree(edge_type, id)
+      # @param [String] edge_type
+      # @param [Teapot::TempestDB::Node] node
+      # @return [Integer]
+      def in_degree(edge_type, node)
         @thrift_client.with_retries { |executor|
-          executor.inDegree(edge_type, id)
+          executor.inDegree(edge_type, node)
         }
       end
 
-      def out_neighbors(edge_type, id)
+      # @param [String] edge_type
+      # @param [Teapot::TempestDB::Node] node
+      # @return [Array<Teapot::TempestDB::Node>]
+      def out_neighbors(edge_type, node)
         @thrift_client.with_retries { |executor|
-          executor.outNeighbors(edge_type, id)
+          executor.outNeighbors(edge_type, node)
         }
       end
 
-      def in_neighbors(edge_type, id)
+      # @param [String] edge_type
+      # @param [Teapot::TempestDB::Node] node
+      # @return [Array<Teapot::TempestDB::Node>]
+      def in_neighbors(edge_type, node)
         @thrift_client.with_retries { |executor|
-          executor.inNeighbors(edge_type, id)
+          executor.inNeighbors(edge_type, node)
         }
       end
 
-      def out_neighbor(edge_type, id, i)
+      # @param [String] edge_type
+      # @param [Teapot::TempestDB::Node] node
+      # @param [Integer] i
+      # @return [Teapot::TempestDB::Node]
+      def out_neighbor(edge_type, node, i)
         @thrift_client.with_retries { |executor|
-          executor.outNeighbor(edge_type, id, i)
+          executor.outNeighbor(edge_type, node, i)
         }
       end
 
-      def in_neighbor(edge_type, id, i)
+      # @param [String] edge_type
+      # @param [Teapot::TempestDB::Node] node
+      # @param [Integer] i
+      # @return [Teapot::TempestDB::Node]
+      def in_neighbor(edge_type, node, i)
         @thrift_client.with_retries { |executor|
-          executor.inNeighbor(edge_type, id, i)
-        }
-      end
-
-      def max_node_id(edge_type)
-        @thrift_client.with_retries { |executor|
-          executor.maxNodeId(edge_type)
+          executor.inNeighbor(edge_type, node, i)
         }
       end
 
+      # @param [String] edge_type
+      # @return [Integer]
       def node_count(edge_type)
         @thrift_client.with_retries { |executor|
           executor.nodeCount(edge_type)
         }
       end
 
+      # @param [String] edge_type
+      # @return [Integer]
       def edge_count(edge_type)
         @thrift_client.with_retries { |executor|
           executor.edgeCount(edge_type)
         }
       end
 
-      def ppr_single_target(edge_type, seed_node_ids, target_node_id, bi_ppr_params)
+      # @param [String] edge_type
+      # @param [Array<Teapot::TempestDB::Node>] seed_nodes
+      # @param [Teapot::TempestDB::Node] target_node
+      # @param [Teapot::TempestDB::BidirectionalPPRParams] bi_ppr_params
+      # @return [Double]
+      def ppr_single_target(edge_type, seed_nodes, target_node, bi_ppr_params)
         @thrift_client.with_retries { |executor|
-          executor.pprSingleTarget(edge_type, seed_node_ids, target_node_id, bi_ppr_params)
+          executor.pprSingleTarget(edge_type, seed_nodes, target_node, bi_ppr_params)
         }
       end
 
-      def ppr(edge_type, seed_node_ids, seed_node_type, target_node_type, mc_ppr_params)
+      # @param [Teapot::TempestDB::Node] source
+      # @param [Array<String>] edge_types
+      # @param [Integer] max_size
+      # @return [Array<Teapot::TempestDB::Node>]
+      def connected_component(source, edge_types, max_size = (1 << 31) - 1)
         @thrift_client.with_retries { |executor|
-          executor.ppr(edge_type, seed_node_ids, seed_node_type, target_node_type, mc_ppr_params)
+          executor.connectedComponent(source, edge_types, max_size)
         }
       end
 
+      # @param [String] node_type
+      # @param [String] sql_clause
+      # @return [Teapot::TempestDB::Node]
       def nodes(node_type, sql_clause)
         @thrift_client.with_retries { |executor|
           executor.nodes(node_type, sql_clause)
         }
       end
 
-      def add_edges(edge_type, ids1, ids2)
+      # @param [Teapot::TempestDB::Node] node
+      def add_node(node)
+        @thrift_client.with_retries { |executor|
+          executor.addNode(node)
+        }
+      end
+
+      # @param [Teapot::TempestDB::Node] node
+      # @param [String] attribute_name
+      # @param [String] attribute_value
+      def set_node_attribute(node, attribute_name, attribute_value)
+        @thrift_client.with_retries { |executor|
+          executor.setNodeAttribute(node, attribute_name, attribute_value)
+        }
+      end
+
+      # @param [String] edge_type
+      # @param [Array<Teapot::TempestDB::Node>] source_nodes
+      # @param [Array<Teapot::TempestDB::Node>] target_nodes
+      def add_edges(edge_type, source_nodes, target_nodes)
         @thrift_client.with_retries { |executor|
-          executor.addEdges(edge_type, ids1, ids2)
+          executor.addEdges(edge_type, source_nodes, target_nodes)
         }
       end
+
+      # @param [String] edge_type
+      # @param [Teapot::TempestDB::Node] source_node
+      # @param [Teapot::TempestDB::Node] target_node
+      def add_edge(edge_type, source_node, target_node)
+        add_edges(edge_type, [source_node], [target_node])
+      end
     end
   end
 end
 
+# @param [String] server
+# @param [Integer] port
+# @return [Teapot::TempestDB::TempestClient]
 def get_tempest_client(server = 'localhost', port = 10001)
   Teapot::TempestDB::TempestClient.new(server, port)
 end
 
+# @return [Hash]
 def get_empty_filter()
   Hash.new
 end
+
+# @param [Teapot::TempestDB::Node] node
+# @return [Array]
+def node_to_pair(node)
+  [node.type, node.id]
+end
+
+# @param [String] type
+# @param [String] id
+# @return [Teapot::TempestDB::Node]
+def make_node(type, id)
+  Teapot::TempestDB::Node.new({'type' => type, 'id' => id})
+end