gizzmo 0.1.4
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- data/.document +5 -0
- data/.gitignore +21 -0
- data/LICENSE +177 -0
- data/README.rdoc +45 -0
- data/Rakefile +58 -0
- data/VERSION +1 -0
- data/bin/gizzmo +3 -0
- data/gizzmo.gemspec +72 -0
- data/lib/gizzard.rb +4 -0
- data/lib/gizzard/commands.rb +195 -0
- data/lib/gizzard/thrift.rb +160 -0
- data/lib/gizzmo.rb +172 -0
- data/lib/vendor/thrift_client/simple.rb +334 -0
- data/test/config.yaml +2 -0
- data/test/expected/dry-wrap-table_b_0.txt +5 -0
- data/test/expected/empty-file.txt +0 -0
- data/test/expected/find-only-sql-shard-type.txt +20 -0
- data/test/expected/help-info.txt +1 -0
- data/test/expected/info.txt +30 -0
- data/test/expected/links-for-replicating_table_b_0.txt +2 -0
- data/test/expected/links-for-table_b_0.txt +1 -0
- data/test/expected/links-for-table_repl_0.txt +2 -0
- data/test/expected/original-find.txt +30 -0
- data/test/expected/subtree.txt +33 -0
- data/test/expected/unwrapped-replicating_table_b_0.txt +1 -0
- data/test/expected/unwrapped-table_b_0.txt +1 -0
- data/test/expected/wrap-table_b_0.txt +1 -0
- data/test/helper.rb +10 -0
- data/test/recreate.sql +32 -0
- data/test/test.sh +60 -0
- metadata +92 -0
@@ -0,0 +1,160 @@
|
|
1
|
+
#!/usr/bin/env ruby
|
2
|
+
require 'vendor/thrift_client/simple'
|
3
|
+
|
4
|
+
module Gizzard
|
5
|
+
module Thrift
|
6
|
+
T = ThriftClient::Simple
|
7
|
+
|
8
|
+
def self.struct(*args)
|
9
|
+
T::StructType.new(*args)
|
10
|
+
end
|
11
|
+
|
12
|
+
ShardException = T.make_exception(:ShardException,
|
13
|
+
T::Field.new(:description, T::STRING, 1)
|
14
|
+
)
|
15
|
+
|
16
|
+
ShardId = T.make_struct(:ShardId,
|
17
|
+
T::Field.new(:hostname, T::STRING, 1),
|
18
|
+
T::Field.new(:table_prefix, T::STRING, 2)
|
19
|
+
)
|
20
|
+
|
21
|
+
class ShardId
|
22
|
+
def inspect
|
23
|
+
"#{hostname}/#{table_prefix}"
|
24
|
+
end
|
25
|
+
|
26
|
+
alias_method :to_unix, :inspect
|
27
|
+
|
28
|
+
def self.parse(string)
|
29
|
+
new(*string.split("/"))
|
30
|
+
end
|
31
|
+
end
|
32
|
+
|
33
|
+
ShardInfo = T.make_struct(:ShardInfo,
|
34
|
+
T::Field.new(:id, struct(ShardId), 1),
|
35
|
+
T::Field.new(:class_name, T::STRING, 2),
|
36
|
+
T::Field.new(:source_type, T::STRING, 3),
|
37
|
+
T::Field.new(:destination_type, T::STRING, 4),
|
38
|
+
T::Field.new(:busy, T::I32, 5)
|
39
|
+
)
|
40
|
+
|
41
|
+
class ShardInfo
|
42
|
+
def busy?
|
43
|
+
busy && busy > 0
|
44
|
+
end
|
45
|
+
|
46
|
+
def inspect(short = false)
|
47
|
+
"#{id.inspect}" + (busy? ? " (BUSY)" : "")
|
48
|
+
end
|
49
|
+
|
50
|
+
def to_unix
|
51
|
+
[id.to_unix, class_name, busy? ? "busy" : "unbusy"].join("\t")
|
52
|
+
end
|
53
|
+
end
|
54
|
+
|
55
|
+
LinkInfo = T.make_struct(:LinkInfo,
|
56
|
+
T::Field.new(:up_id, struct(ShardId), 1),
|
57
|
+
T::Field.new(:down_id, struct(ShardId), 2),
|
58
|
+
T::Field.new(:weight, T::I32, 3)
|
59
|
+
)
|
60
|
+
|
61
|
+
class LinkInfo
|
62
|
+
def inspect
|
63
|
+
"#{up_id.inspect} -> #{down_id.inspect}" + (weight == 1 ? "" : " <#{weight}>")
|
64
|
+
end
|
65
|
+
|
66
|
+
def to_unix
|
67
|
+
[up_id.to_unix, down_id.to_unix, weight].join("\t")
|
68
|
+
end
|
69
|
+
|
70
|
+
end
|
71
|
+
|
72
|
+
ShardMigration = T.make_struct(:ShardMigration,
|
73
|
+
T::Field.new(:source_id, struct(ShardId), 1),
|
74
|
+
T::Field.new(:destination_id, struct(ShardId), 2)
|
75
|
+
)
|
76
|
+
|
77
|
+
Forwarding = T.make_struct(:Forwarding,
|
78
|
+
T::Field.new(:table_id, T::I32, 1),
|
79
|
+
T::Field.new(:base_id, T::I64, 2),
|
80
|
+
T::Field.new(:shard_id, struct(ShardId), 3)
|
81
|
+
)
|
82
|
+
|
83
|
+
class Forwarding
|
84
|
+
#FIXME table_id is not human-readable
|
85
|
+
def inspect
|
86
|
+
"[#{table_id}] #{base_id.to_s(16)} -> #{shard_id.inspect}"
|
87
|
+
end
|
88
|
+
end
|
89
|
+
|
90
|
+
class ShardManager < T::ThriftService
|
91
|
+
def initialize(host, port, log_path, dry_run = false)
|
92
|
+
super(host, port)
|
93
|
+
@dry = dry_run
|
94
|
+
begin
|
95
|
+
@log = File.open(log_path, "a")
|
96
|
+
rescue
|
97
|
+
STDERR.puts "Error opening log file at #{log_path}. Continuing..."
|
98
|
+
end
|
99
|
+
end
|
100
|
+
|
101
|
+
def _proxy(method_name, *args)
|
102
|
+
cls = self.class.ancestors.find { |cls| cls.respond_to?(:_arg_structs) and cls._arg_structs[method_name.to_sym] }
|
103
|
+
arg_class, rv_class = cls._arg_structs[method_name.to_sym]
|
104
|
+
|
105
|
+
# Writing methods return void. Methods should never both read and write. If this assumption
|
106
|
+
# is violated in the future, dry-run will fail!!
|
107
|
+
is_writing_method = rv_class._fields.first.type == ThriftClient::Simple::VOID
|
108
|
+
if @dry && is_writing_method
|
109
|
+
puts "Skipped writing: #{printable(method_name, args)}"
|
110
|
+
else
|
111
|
+
@log.puts printable(method_name, args, true)
|
112
|
+
super(method_name, *args)
|
113
|
+
end
|
114
|
+
rescue ThriftClient::Simple::ThriftException
|
115
|
+
if @dry
|
116
|
+
puts "Skipped reading: #{printable(method_name, args)}"
|
117
|
+
else
|
118
|
+
raise
|
119
|
+
end
|
120
|
+
end
|
121
|
+
|
122
|
+
def printable(method_name, args, timestamp = false)
|
123
|
+
ts = timestamp ? "#{Time.now}\t" : ""
|
124
|
+
"#{ts}#{method_name}(#{args.map{|a| a.inspect}.join(', ')})"
|
125
|
+
end
|
126
|
+
|
127
|
+
|
128
|
+
thrift_method :create_shard, void, field(:shard, struct(ShardInfo), 1), :throws => exception(ShardException)
|
129
|
+
thrift_method :delete_shard, void, field(:id, struct(ShardId), 1)
|
130
|
+
thrift_method :get_shard, struct(ShardInfo), field(:id, struct(ShardId), 1)
|
131
|
+
|
132
|
+
thrift_method :add_link, void, field(:up_id, struct(ShardId), 1), field(:down_id, struct(ShardId), 2), field(:weight, i32, 3)
|
133
|
+
thrift_method :remove_link, void, field(:up_id, struct(ShardId), 1), field(:down_id, struct(ShardId), 2)
|
134
|
+
|
135
|
+
thrift_method :list_upward_links, list(struct(LinkInfo)), field(:id, struct(ShardId), 1)
|
136
|
+
thrift_method :list_downward_links, list(struct(LinkInfo)), field(:id, struct(ShardId), 1)
|
137
|
+
|
138
|
+
thrift_method :get_child_shards_of_class, list(struct(ShardInfo)), field(:parent_id, struct(ShardId), 1), field(:class_name, string, 2)
|
139
|
+
|
140
|
+
thrift_method :mark_shard_busy, void, field(:id, struct(ShardId), 1), field(:busy, i32, 2)
|
141
|
+
thrift_method :copy_shard, void, field(:source_id, struct(ShardId), 1), field(:destination_id, struct(ShardId), 2)
|
142
|
+
|
143
|
+
thrift_method :set_forwarding, void, field(:forwarding, struct(Forwarding), 1)
|
144
|
+
thrift_method :replace_forwarding, void, field(:old_id, struct(ShardId), 1), field(:new_id, struct(ShardId), 2)
|
145
|
+
|
146
|
+
thrift_method :get_forwarding, struct(Forwarding), field(:table_id, i32, 1), field(:base_id, i64, 2)
|
147
|
+
thrift_method :get_forwarding_for_shard, struct(Forwarding), field(:shard_id, struct(ShardId), 1)
|
148
|
+
|
149
|
+
thrift_method :get_forwardings, list(struct(Forwarding))
|
150
|
+
thrift_method :reload_forwardings, void
|
151
|
+
|
152
|
+
thrift_method :find_current_forwarding, struct(ShardInfo), field(:table_id, i32, 1), field(:id, i64, 2)
|
153
|
+
|
154
|
+
thrift_method :shards_for_hostname, list(struct(ShardInfo)), field(:hostname, string, 1)
|
155
|
+
thrift_method :get_busy_shards, list(struct(ShardInfo))
|
156
|
+
|
157
|
+
thrift_method :rebuild_schema, void
|
158
|
+
end
|
159
|
+
end
|
160
|
+
end
|
data/lib/gizzmo.rb
ADDED
@@ -0,0 +1,172 @@
|
|
1
|
+
#!/usr/bin/env ruby
|
2
|
+
$: << File.dirname(__FILE__)
|
3
|
+
class HelpNeededError < RuntimeError; end
|
4
|
+
require "optparse"
|
5
|
+
require "ostruct"
|
6
|
+
require "gizzard"
|
7
|
+
require "yaml"
|
8
|
+
|
9
|
+
ORIGINAL_ARGV = ARGV.dup
|
10
|
+
|
11
|
+
# Container for parsed options
|
12
|
+
global_options = OpenStruct.new
|
13
|
+
subcommand_options = OpenStruct.new
|
14
|
+
|
15
|
+
# Leftover arguments
|
16
|
+
argv = nil
|
17
|
+
|
18
|
+
begin
|
19
|
+
YAML.load_file(File.join(ENV["HOME"], ".gizzmorc")).each do |k, v|
|
20
|
+
global_options.send("#{k}=", v)
|
21
|
+
end
|
22
|
+
rescue Errno::ENOENT
|
23
|
+
# Do nothing...
|
24
|
+
rescue => e
|
25
|
+
abort "Unknown error loading ~/.gizzmorc: #{e.message}"
|
26
|
+
end
|
27
|
+
|
28
|
+
subcommands = {
|
29
|
+
'create' => OptionParser.new do |opts|
|
30
|
+
opts.banner = "Usage: #{$0} create [options] HOST TABLE_PREFIX CLASS_NAME"
|
31
|
+
|
32
|
+
opts.on("-s", "--source-type=TYPE") do |s|
|
33
|
+
subcommand_options.source_type = s
|
34
|
+
end
|
35
|
+
|
36
|
+
opts.on("-d", "--destination-type=TYPE") do |s|
|
37
|
+
subcommand_options.destination_type = s
|
38
|
+
end
|
39
|
+
end,
|
40
|
+
'wrap' => OptionParser.new do |opts|
|
41
|
+
opts.banner = "Usage: #{$0} wrap CLASS_NAME SHARD_ID_TO_WRAP [MORE SHARD_IDS...]"
|
42
|
+
end,
|
43
|
+
'subtree' => OptionParser.new do |opts|
|
44
|
+
opts.banner = "Usage: #{$0} subtree SHARD_ID"
|
45
|
+
end,
|
46
|
+
'delete' => OptionParser.new do |opts|
|
47
|
+
opts.banner = "Usage: #{$0} delete SHARD_ID_TO_DELETE [MORE SHARD_IDS]"
|
48
|
+
end,
|
49
|
+
'unwrap' => OptionParser.new do |opts|
|
50
|
+
opts.banner = "Usage: #{$0} unwrap SHARD_ID_TO_REMOVE [MORE SHARD_IDS]"
|
51
|
+
end,
|
52
|
+
'find' => OptionParser.new do |opts|
|
53
|
+
opts.banner = "Usage: #{$0} find [options]"
|
54
|
+
|
55
|
+
opts.on("-t", "--type=TYPE", "Return only shards of the specified TYPE") do |shard_type|
|
56
|
+
subcommand_options.shard_type = shard_type
|
57
|
+
end
|
58
|
+
|
59
|
+
opts.on("-H", "--host=HOST", "HOST of shard") do |shard_host|
|
60
|
+
subcommand_options.shard_host = shard_host
|
61
|
+
end
|
62
|
+
end,
|
63
|
+
'links' => OptionParser.new do |opts|
|
64
|
+
opts.banner = "Usage: #{$0} links SHARD_ID [MORE SHARD_IDS...]"
|
65
|
+
end,
|
66
|
+
'info' => OptionParser.new do |opts|
|
67
|
+
opts.banner = "Usage: #{$0} info SHARD_ID [MORE SHARD_IDS...]"
|
68
|
+
end,
|
69
|
+
'reload' => OptionParser.new do |opts|
|
70
|
+
opts.banner = "Usage: #{$0} reload"
|
71
|
+
end,
|
72
|
+
'addlink' => OptionParser.new do |opts|
|
73
|
+
opts.banner = "Usage: #{$0} link PARENT_SHARD_ID CHILD_SHARD_ID WEIGHT"
|
74
|
+
end,
|
75
|
+
'unlink' => OptionParser.new do |opts|
|
76
|
+
opts.banner = "Usage: #{$0} unlink PARENT_SHARD_ID CHILD_SHARD_ID"
|
77
|
+
end
|
78
|
+
}
|
79
|
+
|
80
|
+
global = OptionParser.new do |opts|
|
81
|
+
opts.banner = "Usage: #{$0} [global-options] SUBCOMMAND [subcommand-options]"
|
82
|
+
opts.separator ""
|
83
|
+
opts.separator "Subcommands:"
|
84
|
+
subcommands.keys.compact.sort.each do |sc|
|
85
|
+
opts.separator " #{sc}"
|
86
|
+
end
|
87
|
+
opts.separator ""
|
88
|
+
opts.separator "You can type `#{$0} help SUBCOMMAND` for help on a specific subcommand."
|
89
|
+
opts.separator ""
|
90
|
+
opts.separator "Global options:"
|
91
|
+
|
92
|
+
opts.on("-H", "--host=HOSTNAME", "HOSTNAME of remote thrift service") do |host|
|
93
|
+
global_options.host = host
|
94
|
+
end
|
95
|
+
|
96
|
+
opts.on("-P", "--port=PORT", "PORT of remote thrift service") do |port|
|
97
|
+
global_options.port = port
|
98
|
+
end
|
99
|
+
|
100
|
+
opts.on("-D", "--dry-run", "") do |port|
|
101
|
+
global_options.dry = true
|
102
|
+
end
|
103
|
+
|
104
|
+
opts.on("-C", "--config=YAML_FILE", "YAML_FILE of option key/values") do |file|
|
105
|
+
YAML.load(File.open(file)).each do |k, v|
|
106
|
+
global_options.send("#{k}=", v)
|
107
|
+
end
|
108
|
+
end
|
109
|
+
|
110
|
+
opts.on("-L", "--log=LOG_FILE", "Path to LOG_FILE") do |file|
|
111
|
+
global_options.log = file
|
112
|
+
end
|
113
|
+
end
|
114
|
+
|
115
|
+
# Print banner if no args
|
116
|
+
if ARGV.length == 0
|
117
|
+
STDERR.puts global
|
118
|
+
exit 1
|
119
|
+
end
|
120
|
+
|
121
|
+
# This
|
122
|
+
def process_nested_parsers(global, subcommands)
|
123
|
+
begin
|
124
|
+
global.order!(ARGV) do |subcommand_name|
|
125
|
+
# puts args.inspect
|
126
|
+
subcommand = subcommands[subcommand_name]
|
127
|
+
argv = subcommand ? subcommand.parse!(ARGV) : ARGV
|
128
|
+
return subcommand_name, argv
|
129
|
+
end
|
130
|
+
rescue => e
|
131
|
+
STDERR.puts e.message
|
132
|
+
exit 1
|
133
|
+
end
|
134
|
+
end
|
135
|
+
|
136
|
+
|
137
|
+
subcommand_name, argv = process_nested_parsers(global, subcommands)
|
138
|
+
|
139
|
+
# Print help sub-banners
|
140
|
+
if subcommand_name == "help"
|
141
|
+
STDERR.puts subcommands[argv.shift] || global
|
142
|
+
exit 1
|
143
|
+
end
|
144
|
+
|
145
|
+
unless subcommands.include?(subcommand_name)
|
146
|
+
STDERR.puts "Subcommand not found: #{subcommand_name}"
|
147
|
+
exit 1
|
148
|
+
end
|
149
|
+
|
150
|
+
log = global_options.log || "/tmp/gizzmo.log"
|
151
|
+
service = Gizzard::Thrift::ShardManager.new(global_options.host, global_options.port, log, global_options.dry)
|
152
|
+
|
153
|
+
begin
|
154
|
+
Gizzard::Command.run(subcommand_name, service, global_options, argv, subcommand_options)
|
155
|
+
rescue HelpNeededError => e
|
156
|
+
if e.class.name != e.message
|
157
|
+
STDERR.puts("=" * 80)
|
158
|
+
STDERR.puts e.message
|
159
|
+
STDERR.puts("=" * 80)
|
160
|
+
end
|
161
|
+
STDERR.puts subcommands[subcommand_name]
|
162
|
+
exit 1
|
163
|
+
rescue ThriftClient::Simple::ThriftException => e
|
164
|
+
STDERR.puts e.message
|
165
|
+
exit 1
|
166
|
+
rescue Errno::EPIPE
|
167
|
+
# This is just us trying to puts into a closed stdout. For example, if you pipe into
|
168
|
+
# head -1, then this script will keep running after head closes. We don't care, and
|
169
|
+
# seeing the backtrace is annoying!
|
170
|
+
rescue Interrupt
|
171
|
+
exit 1
|
172
|
+
end
|
@@ -0,0 +1,334 @@
|
|
1
|
+
require 'socket'
|
2
|
+
require 'getoptlong'
|
3
|
+
|
4
|
+
class ThriftClient
|
5
|
+
|
6
|
+
# This is a simplified form of thrift, useful for clients only, and not
|
7
|
+
# making any attempt to have good performance. It's intended to be used by
|
8
|
+
# small command-line tools that don't want to install a dozen ruby files.
|
9
|
+
module Simple
|
10
|
+
VERSION_1 = 0x8001
|
11
|
+
|
12
|
+
# message types
|
13
|
+
CALL, REPLY, EXCEPTION = (1..3).to_a
|
14
|
+
|
15
|
+
# value types
|
16
|
+
STOP, VOID, BOOL, BYTE, DOUBLE, _, I16, _, I32, _, I64, STRING, STRUCT, MAP, SET, LIST = (0..15).to_a
|
17
|
+
|
18
|
+
FORMATS = {
|
19
|
+
BYTE => "c",
|
20
|
+
DOUBLE => "G",
|
21
|
+
I16 => "n",
|
22
|
+
I32 => "N",
|
23
|
+
}
|
24
|
+
|
25
|
+
SIZES = {
|
26
|
+
BYTE => 1,
|
27
|
+
DOUBLE => 8,
|
28
|
+
I16 => 2,
|
29
|
+
I32 => 4,
|
30
|
+
}
|
31
|
+
|
32
|
+
module ComplexType
|
33
|
+
module Extends
|
34
|
+
def type_id=(n)
|
35
|
+
@type_id = n
|
36
|
+
end
|
37
|
+
|
38
|
+
def type_id
|
39
|
+
@type_id
|
40
|
+
end
|
41
|
+
end
|
42
|
+
|
43
|
+
module Includes
|
44
|
+
def to_i
|
45
|
+
self.class.type_id
|
46
|
+
end
|
47
|
+
|
48
|
+
def to_s
|
49
|
+
args = self.values.map { |v| self.class.type_id == STRUCT ? v.name : v.to_s }.join(", ")
|
50
|
+
"#{self.class.name}.new(#{args})"
|
51
|
+
end
|
52
|
+
end
|
53
|
+
end
|
54
|
+
|
55
|
+
def self.make_type(type_id, name, *args)
|
56
|
+
klass = Struct.new("STT_#{name}", *args)
|
57
|
+
klass.send(:extend, ComplexType::Extends)
|
58
|
+
klass.send(:include, ComplexType::Includes)
|
59
|
+
klass.type_id = type_id
|
60
|
+
klass
|
61
|
+
end
|
62
|
+
|
63
|
+
ListType = make_type(LIST, "ListType", :element_type)
|
64
|
+
MapType = make_type(MAP, "MapType", :key_type, :value_type)
|
65
|
+
SetType = make_type(SET, "SetType", :element_type)
|
66
|
+
StructType = make_type(STRUCT, "StructType", :struct_class)
|
67
|
+
|
68
|
+
class << self
|
69
|
+
def pack_value(type, value)
|
70
|
+
case type
|
71
|
+
when BOOL
|
72
|
+
[ value ? 1 : 0 ].pack("c")
|
73
|
+
when STRING
|
74
|
+
[ value.size, value ].pack("Na*")
|
75
|
+
when I64
|
76
|
+
[ value >> 32, value & 0xffffffff ].pack("NN")
|
77
|
+
when ListType
|
78
|
+
[ type.element_type.to_i, value.size ].pack("cN") + value.map { |item| pack_value(type.element_type, item) }.join("")
|
79
|
+
when MapType
|
80
|
+
[ type.key_type.to_i, type.value_type.to_i, value.size ].pack("ccN") + value.map { |k, v| pack_value(type.key_type, k) + pack_value(type.value_type, v) }.join("")
|
81
|
+
when SetType
|
82
|
+
[ type.element_type.to_i, value.size ].pack("cN") + value.map { |item| pack_value(type.element_type, item) }.join("")
|
83
|
+
when StructType
|
84
|
+
value._pack
|
85
|
+
else
|
86
|
+
[ value ].pack(FORMATS[type])
|
87
|
+
end
|
88
|
+
end
|
89
|
+
|
90
|
+
def pack_request(method_name, arg_struct, request_id=0)
|
91
|
+
[ VERSION_1, CALL, method_name.to_s.size, method_name.to_s, request_id, arg_struct._pack ].pack("nnNa*Na*")
|
92
|
+
end
|
93
|
+
|
94
|
+
def read_value(s, type)
|
95
|
+
case type
|
96
|
+
when BOOL
|
97
|
+
s.read(1).unpack("c").first != 0
|
98
|
+
when STRING
|
99
|
+
len = s.read(4).unpack("N").first
|
100
|
+
s.read(len)
|
101
|
+
when I64
|
102
|
+
hi, lo = s.read(8).unpack("NN")
|
103
|
+
rv = (hi << 32) | lo
|
104
|
+
(rv >= (1 << 63)) ? (rv - (1 << 64)) : rv
|
105
|
+
when LIST
|
106
|
+
read_list(s)
|
107
|
+
when MAP
|
108
|
+
read_map(s)
|
109
|
+
when STRUCT
|
110
|
+
read_struct(s, UnknownStruct)
|
111
|
+
when ListType
|
112
|
+
read_list(s, type.element_type)
|
113
|
+
when MapType
|
114
|
+
read_map(s, type.key_type, type.value_type)
|
115
|
+
when StructType
|
116
|
+
read_struct(s, type.struct_class)
|
117
|
+
else
|
118
|
+
rv = s.read(SIZES[type]).unpack(FORMATS[type]).first
|
119
|
+
case type
|
120
|
+
when I16
|
121
|
+
(rv >= (1 << 15)) ? (rv - (1 << 16)) : rv
|
122
|
+
when I32
|
123
|
+
(rv >= (1 << 31)) ? (rv - (1 << 32)) : rv
|
124
|
+
else
|
125
|
+
rv
|
126
|
+
end
|
127
|
+
end
|
128
|
+
end
|
129
|
+
|
130
|
+
def read_list(s, element_type=nil)
|
131
|
+
etype, len = s.read(5).unpack("cN")
|
132
|
+
expected_type = (element_type and element_type.to_i == etype.to_i) ? element_type : etype
|
133
|
+
rv = []
|
134
|
+
len.times do
|
135
|
+
rv << read_value(s, expected_type)
|
136
|
+
end
|
137
|
+
rv
|
138
|
+
end
|
139
|
+
|
140
|
+
def read_map(s, key_type=nil, value_type=nil)
|
141
|
+
ktype, vtype, len = s.read(6).unpack("ccN")
|
142
|
+
rv = {}
|
143
|
+
expected_key_type, expected_value_type = if key_type and value_type and key_type.to_i == ktype and value_type.to_i == vtype
|
144
|
+
[ key_type, value_type ]
|
145
|
+
else
|
146
|
+
[ ktype, vtype ]
|
147
|
+
end
|
148
|
+
len.times do
|
149
|
+
key = read_value(s, expected_key_type)
|
150
|
+
value = read_value(s, expected_value_type)
|
151
|
+
rv[key] = value
|
152
|
+
end
|
153
|
+
rv
|
154
|
+
end
|
155
|
+
|
156
|
+
def read_struct(s, struct_class)
|
157
|
+
struct = struct_class.new
|
158
|
+
while true
|
159
|
+
ftype = s.read(1).unpack("c").first
|
160
|
+
return struct if ftype == STOP
|
161
|
+
fid = s.read(2).unpack("n").first
|
162
|
+
|
163
|
+
if field = struct_class._fields.find { |f| (f.fid == fid) and (f.type.to_i == ftype) }
|
164
|
+
struct[field.name] = read_value(s, field.type)
|
165
|
+
else
|
166
|
+
$stderr.puts "Warning: Unknown struct field encountered. (recieved id: #{fid})"
|
167
|
+
raise "Warning: Unknown struct field encountered. (recieved id: #{fid})"
|
168
|
+
read_value(s, ftype)
|
169
|
+
end
|
170
|
+
end
|
171
|
+
end
|
172
|
+
|
173
|
+
def read_response(s, rv_class)
|
174
|
+
version, message_type, method_name_len = s.read(8).unpack("nnN")
|
175
|
+
method_name = s.read(method_name_len)
|
176
|
+
seq_id = s.read(4).unpack("N").first
|
177
|
+
if message_type == EXCEPTION
|
178
|
+
exception = read_struct(s, ExceptionStruct)
|
179
|
+
raise ThriftException, exception.message
|
180
|
+
end
|
181
|
+
response = read_struct(s, rv_class)
|
182
|
+
raise response.ex if response.respond_to?(:ex) and response.ex
|
183
|
+
[ method_name, seq_id, response.rv ]
|
184
|
+
end
|
185
|
+
end
|
186
|
+
|
187
|
+
## ----------------------------------------
|
188
|
+
|
189
|
+
class Field
|
190
|
+
attr_accessor :name, :type, :fid
|
191
|
+
|
192
|
+
def initialize(name, type, fid)
|
193
|
+
@name = name
|
194
|
+
@type = type
|
195
|
+
@fid = fid
|
196
|
+
end
|
197
|
+
|
198
|
+
def pack(value)
|
199
|
+
value.nil? ? "" : [ type.to_i, fid, ThriftClient::Simple.pack_value(type, value) ].pack("cna*")
|
200
|
+
end
|
201
|
+
end
|
202
|
+
|
203
|
+
class ThriftException < RuntimeError
|
204
|
+
def initialize(reason)
|
205
|
+
@reason = reason
|
206
|
+
end
|
207
|
+
|
208
|
+
def to_s
|
209
|
+
"ThriftException(#{@reason.inspect})"
|
210
|
+
end
|
211
|
+
end
|
212
|
+
|
213
|
+
module ThriftStruct
|
214
|
+
module Include
|
215
|
+
def _pack
|
216
|
+
self.class._fields.map { |f| f.pack(self[f.name]) }.join + [ STOP ].pack("c")
|
217
|
+
end
|
218
|
+
end
|
219
|
+
|
220
|
+
module Extend
|
221
|
+
def _fields
|
222
|
+
@fields
|
223
|
+
end
|
224
|
+
|
225
|
+
def _fields=(f)
|
226
|
+
@fields = f
|
227
|
+
end
|
228
|
+
end
|
229
|
+
end
|
230
|
+
|
231
|
+
def self.make_struct(name, *fields)
|
232
|
+
st_name = "ST_#{name.to_s.tr(':', '_')}"
|
233
|
+
if Struct.constants.include?(st_name)
|
234
|
+
warn "#{caller[0]}: Struct::#{st_name} is already defined; returning original class."
|
235
|
+
Struct.const_get(st_name)
|
236
|
+
else
|
237
|
+
names = fields.map { |f| f.name.to_sym }
|
238
|
+
klass = Struct.new(st_name, *names)
|
239
|
+
klass.send(:include, ThriftStruct::Include)
|
240
|
+
klass.send(:extend, ThriftStruct::Extend)
|
241
|
+
klass._fields = fields
|
242
|
+
klass
|
243
|
+
end
|
244
|
+
end
|
245
|
+
|
246
|
+
def self.make_exception(name, *fields)
|
247
|
+
struct_class = self.make_struct(name, *fields)
|
248
|
+
ex_class = Class.new(StandardError)
|
249
|
+
|
250
|
+
(class << struct_class; self end).send(:define_method, :exception_class) { ex_class }
|
251
|
+
(class << ex_class; self end).send(:define_method, :struct_class) { struct_class }
|
252
|
+
|
253
|
+
ex_class.class_eval do
|
254
|
+
attr_reader :struct
|
255
|
+
|
256
|
+
def initialize
|
257
|
+
@struct = self.class.struct_class.new
|
258
|
+
end
|
259
|
+
|
260
|
+
def self._fields
|
261
|
+
struct_class._fields
|
262
|
+
end
|
263
|
+
|
264
|
+
def to_s
|
265
|
+
method = [:message, :description].find {|m| struct.respond_to? m }
|
266
|
+
struct.send method || :to_s
|
267
|
+
end
|
268
|
+
|
269
|
+
alias message to_s
|
270
|
+
|
271
|
+
def method_missing(method, *args)
|
272
|
+
struct.send(method, *args)
|
273
|
+
end
|
274
|
+
end
|
275
|
+
|
276
|
+
ex_class
|
277
|
+
end
|
278
|
+
|
279
|
+
ExceptionStruct = make_struct(:ProtocolException, Field.new(:message, STRING, 1), Field.new(:type, I32, 2))
|
280
|
+
UnknownStruct = make_struct(:Unknown)
|
281
|
+
|
282
|
+
class ThriftService
|
283
|
+
def initialize(host, port)
|
284
|
+
@host = host
|
285
|
+
@port = port
|
286
|
+
end
|
287
|
+
|
288
|
+
def self._arg_structs
|
289
|
+
@_arg_structs = {} if @_arg_structs.nil?
|
290
|
+
@_arg_structs
|
291
|
+
end
|
292
|
+
|
293
|
+
def self.thrift_method(name, rtype, *args)
|
294
|
+
options = args.last.is_a?(Hash) ? args.pop : {}
|
295
|
+
fields = [ ThriftClient::Simple::Field.new(:rv, rtype, 0),
|
296
|
+
(options[:throws] ? ThriftClient::Simple::Field.new(:ex, options[:throws], 1) : nil)
|
297
|
+
].compact
|
298
|
+
|
299
|
+
arg_struct = ThriftClient::Simple.make_struct("Args__#{self.name}__#{name}", *args)
|
300
|
+
rv_struct = ThriftClient::Simple.make_struct("Retval__#{self.name}__#{name}", *fields)
|
301
|
+
|
302
|
+
_arg_structs[name.to_sym] = [ arg_struct, rv_struct ]
|
303
|
+
|
304
|
+
arg_names = args.map { |a| a.name.to_s }.join(", ")
|
305
|
+
class_eval "def #{name}(#{arg_names}); _proxy(:#{name}#{args.size > 0 ? ', ' : ''}#{arg_names}); end"
|
306
|
+
end
|
307
|
+
|
308
|
+
def _proxy(method_name, *args)
|
309
|
+
cls = self.class.ancestors.find { |cls| cls.respond_to?(:_arg_structs) and cls._arg_structs[method_name.to_sym] }
|
310
|
+
arg_class, rv_class = cls._arg_structs[method_name.to_sym]
|
311
|
+
arg_struct = arg_class.new(*args)
|
312
|
+
sock = TCPSocket.new(@host, @port)
|
313
|
+
sock.write(ThriftClient::Simple.pack_request(method_name, arg_struct))
|
314
|
+
rv = ThriftClient::Simple.read_response(sock, rv_class)
|
315
|
+
sock.close
|
316
|
+
rv[2]
|
317
|
+
end
|
318
|
+
|
319
|
+
# convenience. robey is lazy.
|
320
|
+
{ :field => "Field.new",
|
321
|
+
:struct => "StructType.new",
|
322
|
+
:exception => "StructType.new",
|
323
|
+
:list => "ListType.new",
|
324
|
+
:map => "MapType.new",
|
325
|
+
}.each do |new_name, old_name|
|
326
|
+
class_eval "def self.#{new_name}(*args); ThriftClient::Simple::#{old_name}(*args); end"
|
327
|
+
end
|
328
|
+
|
329
|
+
# alias exception struct
|
330
|
+
|
331
|
+
[ :void, :bool, :byte, :double, :i16, :i32, :i64, :string ].each { |sym| class_eval "def self.#{sym}; ThriftClient::Simple::#{sym.to_s.upcase}; end" }
|
332
|
+
end
|
333
|
+
end
|
334
|
+
end
|