feature #788: Add onequota command

src/authm_mad/remotes/quota/onequota

@@ -0,0 +1,158 @@
+#!/usr/bin/env ruby
+
+# -------------------------------------------------------------------------- #
+# Copyright 2002-2011, OpenNebula Project Leads (OpenNebula.org)             #
+#                                                                            #
+# Licensed under the Apache License, Version 2.0 (the "License"); you may    #
+# not use this file except in compliance with the License. You may obtain    #
+# a copy of the License at                                                   #
+#                                                                            #
+# http://www.apache.org/licenses/LICENSE-2.0                                 #
+#                                                                            #
+# Unless required by applicable law or agreed to in writing, software        #
+# distributed under the License is distributed on an "AS IS" BASIS,          #
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.   #
+# See the License for the specific language governing permissions and        #
+# limitations under the License.                                             #
+#--------------------------------------------------------------------------- #
+
+ONE_LOCATION=ENV["ONE_LOCATION"]
+
+if !ONE_LOCATION
+    RUBY_LIB_LOCATION="/usr/lib/one/ruby"
+else
+    RUBY_LIB_LOCATION=ONE_LOCATION+"/lib/ruby"
+end
+
+$: << RUBY_LIB_LOCATION
+$: << RUBY_LIB_LOCATION+"/cli"
+
+require 'command_parser'
+require 'cli_helper'
+require 'one_helper'
+require 'quota'
+
+require 'pp'
+
+QUOTA_KEYS = Quota::DB_QUOTA_SCHEMA.keys
+
+def show_table(values, usage=nil)
+    size = usage ? 12 : 8
+
+    values.sort!{|v1, v2| v1[:uid]<=>v2[:uid] }
+
+    table=CLIHelper::ShowTable.new(nil, self) do
+        column :uid, "ONE identifier for the User", :size=>4 do |d|
+            d[:uid]
+        end
+
+        QUOTA_KEYS.each { |key|
+            column key, "#{key} quota", :size=>size do |d|
+                if usage
+                    "#{usage[key].to_i}/#{d[key].to_i}"
+                else
+                    "#{d[key].to_i}"
+                end
+            end
+        }
+
+        default :uid, *QUOTA_KEYS
+    end
+
+    table.show(values)
+end
+
+cmd=CommandParser::CmdParser.new(ARGV) do
+    usage "`onequota` <command> [<args>] [<options>]"
+    version OpenNebulaHelper::ONE_VERSION
+
+    quota = Quota.new
+
+    ########################################################################
+    # Global Options
+    ########################################################################
+    set :option, CommandParser::OPTIONS
+
+    ########################################################################
+    # Argument Formats
+    ########################################################################
+    quota_list_desc = <<-EOT.unindent
+        List of quota keys, comma separated.
+        Available quotas: #{QUOTA_KEYS.join(', ')}
+    EOT
+
+    set :format, :quota_list, quota_list_desc do |arg|
+        arg_list = arg.split(',')
+
+        rc = nil
+        arg_list.each do |elem|
+          if !QUOTA_KEYS.include?(elem.to_sym)
+            rc = -1, "#{elem} is not a valid quota"
+          end
+        end
+
+        rc ? rc : [0, arg_list]
+    end
+
+    set :format, :value_list, "List of quota values, comma separated." do |arg|
+        arg_list = arg.split(',')
+        arg_list.map! {|a| a.to_i }
+        [0, arg_list]
+    end
+    
+    set :format, :userid, OpenNebulaHelper.rname_to_id_desc("USER") do |arg|
+        OpenNebulaHelper.rname_to_id(arg, "USER")
+    end
+
+    ########################################################################
+    # Commands
+    ########################################################################
+    set_desc = <<-EOT.unindent
+        Set a quota for a given user.
+        Examples:
+          onequota set 3 cpu 12
+          onequota set 4 cpu,memory,storage 8,4096,10000
+    EOT
+
+    command :set, set_desc, :userid, :quota_list, :value_list do
+        user_id, keys, values = args
+
+        if keys.length != values.length
+            exit_with_code -1, "The number of keys and values does not match"
+        end
+
+        values_hash = Hash.new
+        keys.each_with_index { |k,i|
+            values_hash[k.to_sym] = values[i]
+        }
+
+        quota.set_quota(user_id, values_hash)
+        exit_with_code 0
+    end
+
+    ########################################################################
+    show_desc = "Show the user's quota and usage. (usage/quota)"
+    
+    FORCE={
+        :name  => "force",
+        :short => "-f",
+        :large => "--force",
+        :description => "Force the usage calculation instead of using the cache"
+    }
+
+    command :show, show_desc, :userid, :options=>[FORCE] do
+        user_usage = quota.get_usage(args[0],nil,options[:force])
+        user_quota = quota.get_quota(args[0])
+
+        show_table([user_quota], user_usage)
+        exit_with_code 0
+    end
+
+    ########################################################################
+    list_desc = "List the defined quotas for all the users"
+
+    command :list, list_desc do
+        show_table(quota.get_quota)
+        exit_with_code 0
+    end
+end