github.com/makyo/juju@v0.0.0-20160425123129-2608902037e9/state/multienv.go (about)

     1  // Copyright 2015 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package state
     5  
     6  import (
     7  	"strings"
     8  
     9  	"github.com/juju/errors"
    10  	"gopkg.in/mgo.v2/bson"
    11  )
    12  
    13  // This file contains utility functions related to documents and
    14  // collections that contain data for multiple models.
    15  
    16  // ensureModelUUID returns an model UUID prefixed document ID. The
    17  // prefix is only added if it isn't already there.
    18  func ensureModelUUID(modelUUID, id string) string {
    19  	prefix := modelUUID + ":"
    20  	if strings.HasPrefix(id, prefix) {
    21  		return id
    22  	}
    23  	return prefix + id
    24  }
    25  
    26  // ensureModelUUIDIfString will call ensureModelUUID, but only if the id
    27  // is a string. The id will be left untouched otherwise.
    28  func ensureModelUUIDIfString(modelUUID string, id interface{}) interface{} {
    29  	if id, ok := id.(string); ok {
    30  		return ensureModelUUID(modelUUID, id)
    31  	}
    32  	return id
    33  }
    34  
    35  // splitDocID returns the 2 parts of model UUID prefixed
    36  // document ID. If the id is not in the expected format the final
    37  // return value will be false.
    38  func splitDocID(id string) (string, string, bool) {
    39  	parts := strings.SplitN(id, ":", 2)
    40  	if len(parts) != 2 {
    41  		return "", "", false
    42  	}
    43  	return parts[0], parts[1], true
    44  }
    45  
    46  const modelUUIDRequired = 1
    47  const noModelUUIDInInput = 2
    48  
    49  // mungeDocForMultiEnv takes the value of an txn.Op Insert or $set
    50  // Update and modifies it to be multi-model safe, returning the
    51  // modified document.
    52  func mungeDocForMultiEnv(doc interface{}, modelUUID string, modelUUIDFlags int) (bson.D, error) {
    53  	var bDoc bson.D
    54  	var err error
    55  	if doc != nil {
    56  		bDoc, err = toBsonD(doc)
    57  		if err != nil {
    58  			return nil, errors.Trace(err)
    59  		}
    60  	}
    61  
    62  	modelUUIDSeen := false
    63  	for i, elem := range bDoc {
    64  		switch elem.Name {
    65  		case "_id":
    66  			if id, ok := elem.Value.(string); ok {
    67  				bDoc[i].Value = ensureModelUUID(modelUUID, id)
    68  			} else if subquery, ok := elem.Value.(bson.D); ok {
    69  				munged, err := mungeIDSubQueryForMultiEnv(subquery, modelUUID)
    70  				if err != nil {
    71  					return nil, errors.Trace(err)
    72  				}
    73  				bDoc[i].Value = munged
    74  			}
    75  		case "model-uuid":
    76  			if modelUUIDFlags&noModelUUIDInInput > 0 {
    77  				return nil, errors.New("model-uuid is added automatically and should not be provided")
    78  			}
    79  			modelUUIDSeen = true
    80  			if elem.Value == "" {
    81  				bDoc[i].Value = modelUUID
    82  			} else if elem.Value != modelUUID {
    83  				return nil, errors.Errorf(`bad "model-uuid" value: expected %s, got %s`, modelUUID, elem.Value)
    84  			}
    85  		}
    86  	}
    87  	if modelUUIDFlags&modelUUIDRequired > 0 && !modelUUIDSeen {
    88  		bDoc = append(bDoc, bson.DocElem{"model-uuid", modelUUID})
    89  	}
    90  	return bDoc, nil
    91  }
    92  
    93  func mungeIDSubQueryForMultiEnv(doc interface{}, modelUUID string) (bson.D, error) {
    94  	var bDoc bson.D
    95  	var err error
    96  	if doc != nil {
    97  		bDoc, err = toBsonD(doc)
    98  		if err != nil {
    99  			return nil, errors.Trace(err)
   100  		}
   101  	}
   102  
   103  	for i, elem := range bDoc {
   104  		switch elem.Name {
   105  		case "$in":
   106  			var ids []string
   107  			switch values := elem.Value.(type) {
   108  			case []string:
   109  				ids = values
   110  			case []interface{}:
   111  				for _, value := range values {
   112  					id, ok := value.(string)
   113  					if !ok {
   114  						continue
   115  					}
   116  					ids = append(ids, id)
   117  				}
   118  				if len(ids) != len(values) {
   119  					// We expect the type to be consistently string, so...
   120  					continue
   121  				}
   122  			default:
   123  				continue
   124  			}
   125  
   126  			var fullIDs []string
   127  			for _, id := range ids {
   128  				fullID := ensureModelUUID(modelUUID, id)
   129  				fullIDs = append(fullIDs, fullID)
   130  			}
   131  			bDoc[i].Value = fullIDs
   132  		}
   133  	}
   134  	return bDoc, nil
   135  }
   136  
   137  // toBsonD converts an arbitrary value to a bson.D via marshaling
   138  // through BSON. This is still done even if the input is already a
   139  // bson.D so that we end up with a copy of the input.
   140  func toBsonD(doc interface{}) (bson.D, error) {
   141  	bytes, err := bson.Marshal(doc)
   142  	if err != nil {
   143  		return nil, errors.Annotate(err, "bson marshaling failed")
   144  	}
   145  	var out bson.D
   146  	err = bson.Unmarshal(bytes, &out)
   147  	if err != nil {
   148  		return nil, errors.Annotate(err, "bson unmarshaling failed")
   149  	}
   150  	return out, nil
   151  }