From 5e95592fc5d139b84c0c57a88018e0b93b78e3db Mon Sep 17 00:00:00 2001
From: Yehuda Sadeh <yehuda@redhat.com>
Date: Mon, 7 Jan 2019 10:27:45 -0800
Subject: [PATCH] ragweed: support for storage classes in zone config

Signed-off-by: Yehuda Sadeh <yehuda@redhat.com>
---
 ragweed/framework.py | 62 +++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 59 insertions(+), 3 deletions(-)

diff --git a/ragweed/framework.py b/ragweed/framework.py
index 761c162..9956057 100644
--- a/ragweed/framework.py
+++ b/ragweed/framework.py
@@ -134,6 +134,17 @@ def rtest_decode_json(d):
         return pickle.loads(str(d['__pickle']))
     return d
 
+class RPlacementRule:
+    def __init__(self, rule):
+        r = rule.split('/', 1)
+
+        self.placement_id = r[0]
+
+        if (len(r) == 2):
+            self.storage_class=r[1]
+        else:
+            self.storage_class = 'STANDARD'
+
 
 class RBucket:
     def __init__(self, zone, bucket, bucket_info):
@@ -142,6 +153,12 @@ class RBucket:
         self.name = bucket.name
         self.bucket_info = bucket_info
 
+        try:
+            self.placement_rule = RPlacementRule(self.bucket_info.placement_rule)
+            self.placement_target = self.zone.get_placement_target(self.bucket_info.placement_rule)
+        except:
+            pass
+
     def get_data_pool(self):
         try:
             # old style explicit pool
@@ -151,7 +168,9 @@ class RBucket:
             explicit_pool = self.bucket_info.bucket.explicit_placement.data_pool
         if explicit_pool is not None and explicit_pool != '':
             return explicit_pool
-        return self.zone.get_placement_target(self.bucket_info.placement_rule).data_pool
+
+        return self.placement_target.get_data_pool(self.placement_rule)
+
 
     def get_tail_pool(self, obj_layout):
         try:
@@ -171,7 +190,44 @@ class RBucket:
                 except:
                     pass
 
-        return self.zone.get_placement_target(placement_rule).data_pool
+        pr = RPlacementRule(placement_rule)
+
+        return self.placement_target.get_data_pool(pr)
+
+class RStorageClasses:
+    def __init__(self, config):
+        if hasattr(config, 'storage_classes'):
+            self.storage_classes = config.storage_classes
+        else:
+            try:
+                self.storage_classes = bunch.bunchify({ 'STANDARD': { 'data_pool': config.data_pool }})
+            except:
+                self.storage_classes = None
+                pass
+
+    def get(self, storage_class):
+        assert(self.storage_classes != None)
+        try:
+            if not storage_class:
+                storage_class = 'STANDARD'
+            sc = self.storage_classes[storage_class]
+        except:
+            eq('could not find storage class ' + storage_class, 0)
+
+        return sc
+
+class RPlacementTarget:
+    def __init__(self, name, config):
+        self.name = name
+        self.index_pool = config.index_pool
+        self.data_extra_pool = config.data_extra_pool
+        self.storage_classes = RStorageClasses(config)
+
+        if not self.data_extra_pool:
+            self.data_extra_pool = self.storage_classes.get_data_pool('STANDARD')
+
+    def get_data_pool(self, placement_rule):
+        return self.storage_classes.get(placement_rule.storage_class).data_pool
 
 class RZone:
     def __init__(self, conn):
@@ -194,7 +250,7 @@ class RZone:
             plid = self.zone_params.default_placement
 
         try:
-            return self.placement_targets[plid]
+            return RPlacementTarget(plid, self.placement_targets[plid])
         except:
             pass
 
-- 
2.47.3