<commit_before>package koding

import (
	"errors"
	"koding/db/models"
	"koding/db/mongodb"
	"koding/kites/kloud/protocol"
	"time"

	"github.com/koding/logging"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

const (
	domainCollection = "jDomainAliases"
)

// DomainDocument defines a single MongoDB document in the jDomains collection
type DomainDocument struct {
	Id         bson.ObjectId `bson:"_id" json:"-"`
	OriginId   bson.ObjectId `bson:"originId"`
	MachineId  bson.ObjectId `bson:"machineId"`
	DomainName string        `bson:"domain"`
	CreatedAt  time.Time     `bson:"createdAt"`
	ModifiedAt time.Time     `bson:"modifiedAt"`
}

type Domains struct {
	DB  *mongodb.MongoDB
	Log logging.Logger
}

func NewDomainStorage(db *mongodb.MongoDB) *Domains {
	return &Domains{
		DB:  db,
		Log: logging.NewLogger("kloud-domain"),
	}
}

func (d *Domains) Add(domain *protocol.Domain) error {
	var account *models.Account
	if err := d.DB.Run("jAccounts", func(c *mgo.Collection) error {
		return c.Find(bson.M{"profile.nickname": domain.Username}).One(&account)
	}); err != nil {
		d.Log.Error("Could not fetch account %v: err: %v", domain.Username, err)
		return errors.New("could not fetch account from DB")
	}

	doc := &DomainDocument{
		Id:         bson.NewObjectId(),
		OriginId:   account.Id,
		MachineId:  bson.ObjectIdHex(domain.MachineId),
		DomainName: domain.Name,
		CreatedAt:  time.Now().UTC(),
		ModifiedAt: time.Now().UTC(),
	}

	err := d.DB.Run(domainCollection, func(c *mgo.Collection) error {
		_, err := c.Upsert(bson.M{"domain": domain.Name}, doc)
		return err
	})

	if err != nil {
		d.Log.Error("Could not add %v: err: %v", doc, err)
		return errors.New("could not add account from DB")
	}

	return nil
}

func (d *Domains) Delete(name string) error {
	err := d.DB.Run(domainCollection, func(c *mgo.Collection) error {
		return c.Remove(bson.M{"domain": name})
	})

	if err != nil {
		d.Log.Error("Could not delete %v: err: %v", name, err)
		return errors.New("could not delete domain from DB")
	}

	return nil
}

func (d *Domains) Get(name string) (*protocol.Domain, error) {
	doc := &DomainDocument{}
	err := d.DB.Run(domainCollection, func(c *mgo.Collection) error {
		return c.Find(bson.M{"domain": name}).One(&doc)
	})
	if err != nil {
		return nil, err
	}

	return &protocol.Domain{
		MachineId: doc.MachineId.Hex(),
		Name:      doc.DomainName,
	}, nil
}

func (d *Domains) UpdateMachine(name, machineId string) error {
	err := d.DB.Run(domainCollection, func(c *mgo.Collection) error {
		return c.Update(bson.M{"domain": name},
			bson.M{"$set": bson.M{
				"machineId":  machineId,
				"modifiedAt": time.Now().UTC(),
			}})
	})

	if err != nil {
		d.Log.Error("Could not update %v: err: %v", name, err)
		return errors.New("could not update domain from DB")
	}

	return nil
}

// UpdateDomain sets the ip to the given domain. If there is no record a new
// record will be created otherwise existing record is updated. This is just a
// helper method that uses our DNS struct.
func (p *Provider) UpdateDomain(ip, domain, username string) error {
	if err := p.DNS.Validate(domain, username); err != nil {
		return err
	}

	// Check if the record exist, if yes update the ip instead of creating a new one.
	record, err := p.DNS.Get(domain)
	if err == ErrNoRecord {
		if err := p.DNS.Create(domain, ip); err != nil {
			return err
		}
	} else if err != nil {
		// If it's something else just return it
		return err
	}

	// Means the record exist, update it
	if err == nil {
		p.Log.Warning("Domain '%s' already exists (that shouldn't happen). Going to update to new IP", domain)
		if err := p.DNS.Update(domain, record.IP, ip); err != nil {
			return err
		}
	}

	return nil
}

// userDomains returns a list of all domains that are bound to the given
// machineId
func (p *Provider) userDomains(machineId string) ([]DomainDocument, error) {
	domains := make([]DomainDocument, 0)

	query := func(c *mgo.Collection) error {
		domain := DomainDocument{}

		iter := c.Find(bson.M{"machineId": bson.ObjectIdHex(machineId)}).Batch(20).Iter()
		for iter.Next(&domain) {
			domains = append(domains, domain)
		}

		return iter.Close()
	}

	if err := p.Session.Run(domainCollection, query); err != nil {
		return nil, err
	}

	return domains, nil
}
<commit_msg>kloud/domain: fix UpdateMachine method<commit_after>package koding

import (
	"errors"
	"fmt"
	"koding/db/models"
	"koding/db/mongodb"
	"koding/kites/kloud/protocol"
	"time"

	"github.com/koding/logging"

	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
)

const (
	domainCollection = "jDomainAliases"
)

// DomainDocument defines a single MongoDB document in the jDomains collection
type DomainDocument struct {
	Id         bson.ObjectId `bson:"_id" json:"-"`
	OriginId   bson.ObjectId `bson:"originId"`
	MachineId  bson.ObjectId `bson:"machineId"`
	DomainName string        `bson:"domain"`
	CreatedAt  time.Time     `bson:"createdAt"`
	ModifiedAt time.Time     `bson:"modifiedAt"`
}

type Domains struct {
	DB  *mongodb.MongoDB
	Log logging.Logger
}

func NewDomainStorage(db *mongodb.MongoDB) *Domains {
	return &Domains{
		DB:  db,
		Log: logging.NewLogger("kloud-domain"),
	}
}

func (d *Domains) Add(domain *protocol.Domain) error {
	var account *models.Account
	if err := d.DB.Run("jAccounts", func(c *mgo.Collection) error {
		return c.Find(bson.M{"profile.nickname": domain.Username}).One(&account)
	}); err != nil {
		d.Log.Error("Could not fetch account %v: err: %v", domain.Username, err)
		return errors.New("could not fetch account from DB")
	}

	doc := &DomainDocument{
		Id:         bson.NewObjectId(),
		OriginId:   account.Id,
		MachineId:  bson.ObjectIdHex(domain.MachineId),
		DomainName: domain.Name,
		CreatedAt:  time.Now().UTC(),
		ModifiedAt: time.Now().UTC(),
	}

	err := d.DB.Run(domainCollection, func(c *mgo.Collection) error {
		_, err := c.Upsert(bson.M{"domain": domain.Name}, doc)
		return err
	})

	if err != nil {
		d.Log.Error("Could not add %v: err: %v", doc, err)
		return errors.New("could not add account from DB")
	}

	return nil
}

func (d *Domains) Delete(name string) error {
	err := d.DB.Run(domainCollection, func(c *mgo.Collection) error {
		return c.Remove(bson.M{"domain": name})
	})

	if err != nil {
		d.Log.Error("Could not delete %v: err: %v", name, err)
		return errors.New("could not delete domain from DB")
	}

	return nil
}

func (d *Domains) Get(name string) (*protocol.Domain, error) {
	doc := &DomainDocument{}
	err := d.DB.Run(domainCollection, func(c *mgo.Collection) error {
		return c.Find(bson.M{"domain": name}).One(&doc)
	})
	if err != nil {
		return nil, err
	}

	return &protocol.Domain{
		MachineId: doc.MachineId.Hex(),
		Name:      doc.DomainName,
	}, nil
}

func (d *Domains) UpdateMachine(name, machineId string) error {
	updateData := bson.M{
		"machineId":  "",
		"modifiedAt": time.Now().UTC(),
	}

	if machineId != "" {
		if !bson.IsObjectIdHex(machineId) {
			return fmt.Errorf("'%s' is not a valid object Id")
		}

		updateData["machineId"] = bson.ObjectIdHex(machineId)
	}

	err := d.DB.Run(domainCollection, func(c *mgo.Collection) error {
		return c.Update(bson.M{"domain": name},
			bson.M{"$set": updateData},
		)
	})

	if err != nil {
		d.Log.Error("Could not update %v: err: %v", name, err)
		return errors.New("could not update domain from DB")
	}

	return nil
}

// UpdateDomain sets the ip to the given domain. If there is no record a new
// record will be created otherwise existing record is updated. This is just a
// helper method that uses our DNS struct.
func (p *Provider) UpdateDomain(ip, domain, username string) error {
	if err := p.DNS.Validate(domain, username); err != nil {
		return err
	}

	// Check if the record exist, if yes update the ip instead of creating a new one.
	record, err := p.DNS.Get(domain)
	if err == ErrNoRecord {
		if err := p.DNS.Create(domain, ip); err != nil {
			return err
		}
	} else if err != nil {
		// If it's something else just return it
		return err
	}

	// Means the record exist, update it
	if err == nil {
		p.Log.Warning("Domain '%s' already exists (that shouldn't happen). Going to update to new IP", domain)
		if err := p.DNS.Update(domain, record.IP, ip); err != nil {
			return err
		}
	}

	return nil
}

// userDomains returns a list of all domains that are bound to the given
// machineId
func (p *Provider) userDomains(machineId string) ([]DomainDocument, error) {
	domains := make([]DomainDocument, 0)

	query := func(c *mgo.Collection) error {
		domain := DomainDocument{}

		iter := c.Find(bson.M{"machineId": bson.ObjectIdHex(machineId)}).Batch(20).Iter()
		for iter.Next(&domain) {
			domains = append(domains, domain)
		}

		return iter.Close()
	}

	if err := p.Session.Run(domainCollection, query); err != nil {
		return nil, err
	}

	return domains, nil
}

{'fraction_non_alphanumeric': 0.14763801935116677, 'fraction_numerical': 0.0011383039271485487, 'mean_word_length': 3.7034261241970023, 'pattern_counts': {'":': 18, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5989669', 'n_tokens_mistral': 3267, 'n_tokens_neox': 3088, 'n_words': 877}