Golang db.Conn函数代码示例

func (s segregatedScheduler) Schedule(opts docker.CreateContainerOptions) (string, *docker.Container, error) {
	conn, err := db.Conn()
	if err != nil {
		return "", nil, err
	}
	defer conn.Close()
	var cont container
	coll := collection()
	defer coll.Close()
	err = coll.Find(bson.M{"name": opts.Name}).One(&cont)
	if err != nil {
		return "", nil, err
	}
	app, err := app.GetByName(cont.AppName)
	if err != nil {
		return s.fallback(opts)
	}
	var nodes []node
	query := bson.M{"teams": bson.M{"$in": app.Teams}}
	err = conn.Collection(schedulerCollection).Find(query).All(&nodes)
	if err != nil || len(nodes) < 1 {
		return s.fallback(opts)
	}
	return s.handle(opts, nodes)
}

func removeUserFromTeam(w http.ResponseWriter, r *http.Request, t *auth.Token) error {
	email := r.URL.Query().Get(":user")
	teamName := r.URL.Query().Get(":team")
	u, err := t.User()
	if err != nil {
		return err
	}
	rec.Log(u.Email, "remove-user-from-team", "team="+teamName, "user="+email)
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	team, err := auth.GetTeam(teamName)
	if err != nil {
		return &errors.HTTP{Code: http.StatusNotFound, Message: "Team not found"}
	}
	if !team.ContainsUser(u) {
		msg := fmt.Sprintf("You are not authorized to remove a member from the team %s", team.Name)
		return &errors.HTTP{Code: http.StatusUnauthorized, Message: msg}
	}
	if len(team.Users) == 1 {
		msg := "You can not remove this user from this team, because it is the last user within the team, and a team can not be orphaned"
		return &errors.HTTP{Code: http.StatusForbidden, Message: msg}
	}
	user, err := auth.GetUserByEmail(email)
	if err != nil {
		return &errors.HTTP{Code: http.StatusNotFound, Message: err.Error()}
	}
	err = removeUserFromTeamInGandalf(user, team.Name)
	if err != nil {
		return nil
	}
	return removeUserFromTeamInDatabase(user, team)
}

func (s *WriterSuite) SetUpSuite(c *gocheck.C) {
	var err error
	config.Set("database:url", "127.0.0.1:27017")
	config.Set("database:name", "tsuru_api_writer_test")
	s.conn, err = db.Conn()
	c.Assert(err, gocheck.IsNil)
}

func update(units []provision.Unit) {
	log.Print("updating status from provisioner")
	var l AppList
	for _, unit := range units {
		a, index := l.Search(unit.AppName)
		if index > -1 {
			err := a.Get()
			if err != nil {
				log.Printf("collector: app %q not found. Skipping.\n", unit.AppName)
				continue
			}
		}
		u := app.Unit{}
		u.Name = unit.Name
		u.Type = unit.Type
		u.Machine = unit.Machine
		u.InstanceId = unit.InstanceId
		u.Ip = unit.Ip
		u.State = string(unit.Status)
		a.AddUnit(&u)
		if index > -1 {
			l.Add(a, index)
		}
	}
	conn, err := db.Conn()
	if err != nil {
		log.Printf("collector failed to connect to the database: %s", err)
		return
	}
	defer conn.Close()
	for _, a := range l {
		a.Ip, _ = app.Provisioner.Addr(a)
		conn.Apps().Update(bson.M{"name": a.Name}, a)
	}
}

func collection() (*storage.Collection, error) {
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	return conn.Collection("routers"), nil
}

func (s *S) TestUserStartPasswordReset(c *gocheck.C) {
	conn, err := db.Conn()
	c.Assert(err, gocheck.IsNil)
	defer conn.Close()
	defer s.server.Reset()
	u := User{Email: "[email protected]", Password: "123456"}
	err = u.StartPasswordReset()
	c.Assert(err, gocheck.IsNil)
	defer conn.PasswordTokens().Remove(bson.M{"useremail": u.Email})
	var token passwordToken
	err = conn.PasswordTokens().Find(bson.M{"useremail": u.Email}).One(&token)
	c.Assert(err, gocheck.IsNil)
	time.Sleep(1e9) // Let the email flow.
	s.server.Lock()
	defer s.server.Unlock()
	c.Assert(s.server.MailBox, gocheck.HasLen, 1)
	m := s.server.MailBox[0]
	c.Assert(m.From, gocheck.Equals, "root")
	c.Assert(m.To, gocheck.DeepEquals, []string{u.Email})
	var buf bytes.Buffer
	err = resetEmailData.Execute(&buf, token)
	c.Assert(err, gocheck.IsNil)
	expected := strings.Replace(buf.String(), "\n", "\r\n", -1) + "\r\n"
	c.Assert(string(m.Data), gocheck.Equals, expected)
}

func addUserToTeam(email, teamName string, u *auth.User) error {
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	team, user := new(auth.Team), new(auth.User)
	selector := bson.M{"_id": teamName}
	if err := conn.Teams().Find(selector).One(team); err != nil {
		return &errors.Http{Code: http.StatusNotFound, Message: "Team not found"}
	}
	if !team.ContainsUser(u) {
		msg := fmt.Sprintf("You are not authorized to add new users to the team %s", team.Name)
		return &errors.Http{Code: http.StatusUnauthorized, Message: msg}
	}
	if err := conn.Users().Find(bson.M{"email": email}).One(user); err != nil {
		return &errors.Http{Code: http.StatusNotFound, Message: "User not found"}
	}
	actions := []*action.Action{
		&addUserToTeamInGandalfAction,
		&addUserToTeamInDatabaseAction,
	}
	pipeline := action.NewPipeline(actions...)
	return pipeline.Execute(user.Email, u, team)
}

func removeOldTokens(userEmail string) error {
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	var limit int
	if limit, err = config.GetInt("auth:max-simultaneous-sessions"); err != nil {
		return err
	}
	count, err := conn.Tokens().Find(bson.M{"useremail": userEmail}).Count()
	if err != nil {
		return err
	}
	diff := count - limit
	if diff < 1 {
		return nil
	}
	var tokens []map[string]interface{}
	err = conn.Tokens().Find(bson.M{"useremail": userEmail}).Select(bson.M{"_id": 1}).Limit(diff).All(&tokens)
	if err != nil {
		return nil
	}
	ids := make([]interface{}, 0, len(tokens))
	for _, token := range tokens {
		ids = append(ids, token["_id"])
	}
	_, err = conn.Tokens().RemoveAll(bson.M{"_id": bson.M{"$in": ids}})
	return err
}

func (s *S) TestReserveAppIsSafe(c *gocheck.C) {
	defer runtime.GOMAXPROCS(runtime.GOMAXPROCS(runtime.NumCPU()))
	email := "[email protected]"
	user := &User{
		Email: email, Password: "123456",
		Quota: quota.Quota{Limit: 10, InUse: 0},
	}
	err := user.Create()
	c.Assert(err, gocheck.IsNil)
	conn, err := db.Conn()
	c.Assert(err, gocheck.IsNil)
	defer conn.Close()
	defer conn.Users().Remove(bson.M{"email": user.Email})
	var wg sync.WaitGroup
	for i := 0; i < 24; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			ReserveApp(user)
		}()
	}
	wg.Wait()
	user, err = GetUserByEmail(email)
	c.Assert(err, gocheck.IsNil)
	c.Assert(user.Quota.InUse, gocheck.Equals, 10)
}

func (s *S) TestGetImageAppWhenDeployIsMultipleOf10(c *gocheck.C) {
	var request http.Request
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		request = *r
	}))
	defer server.Close()
	u, _ := url.Parse(server.URL)
	imageRepo := u.Host + "/tsuru/python"
	err := newImage(imageRepo, s.server.URL())
	c.Assert(err, gocheck.IsNil)
	conn, err := db.Conn()
	c.Assert(err, gocheck.IsNil)
	defer conn.Close()
	units := []app.Unit{{Name: "4fa6e0f0c678"}, {Name: "e90e34656806"}}
	app := &app.App{Name: "app1", Platform: "python", Deploys: 20, Units: units}
	err = conn.Apps().Insert(app)
	c.Assert(err, gocheck.IsNil)
	defer conn.Apps().Remove(bson.M{"name": "app1"})
	cont := container{ID: "bleble", Type: "python", AppName: "app1", Image: imageRepo}
	coll := collection()
	err = coll.Insert(cont)
	c.Assert(err, gocheck.IsNil)
	defer coll.Close()
	c.Assert(err, gocheck.IsNil)
	defer coll.RemoveAll(bson.M{"id": "bleble"})
	img := getImage(app)
	repoNamespace, err := config.GetString("docker:repository-namespace")
	c.Assert(err, gocheck.IsNil)
	c.Assert(img, gocheck.Equals, fmt.Sprintf("%s/python", repoNamespace))
	c.Assert(request.Method, gocheck.Equals, "DELETE")
	path := "/v1/repositories/tsuru/python/tags"
	c.Assert(request.URL.Path, gocheck.Equals, path)
}

func (s *S) TestUsePlatformImage(c *gocheck.C) {
	conn, err := db.Conn()
	c.Assert(err, gocheck.IsNil)
	defer conn.Close()
	units := []app.Unit{{Name: "4fa6e0f0c678"}, {Name: "e90e34656806"}}
	app1 := &app.App{Name: "app1", Platform: "python", Deploys: 40, Units: units}
	err = conn.Apps().Insert(app1)
	c.Assert(err, gocheck.IsNil)
	ok := usePlatformImage(app1)
	c.Assert(ok, gocheck.Equals, true)
	defer conn.Apps().Remove(bson.M{"name": "app1"})
	app2 := &app.App{Name: "app2", Platform: "python", Deploys: 20, Units: units}
	err = conn.Apps().Insert(app2)
	c.Assert(err, gocheck.IsNil)
	ok = usePlatformImage(app2)
	c.Assert(ok, gocheck.Equals, true)
	defer conn.Apps().Remove(bson.M{"name": "app2"})
	app3 := &app.App{Name: "app3", Platform: "python", Deploys: 0, Units: units}
	err = conn.Apps().Insert(app3)
	c.Assert(err, gocheck.IsNil)
	ok = usePlatformImage(app3)
	c.Assert(ok, gocheck.Equals, false)
	defer conn.Apps().Remove(bson.M{"name": "app3"})
	app4 := &app.App{Name: "app4", Platform: "python", Deploys: 19, Units: units}
	err = conn.Apps().Insert(app4)
	c.Assert(err, gocheck.IsNil)
	ok = usePlatformImage(app4)
	c.Assert(ok, gocheck.Equals, false)
	defer conn.Apps().Remove(bson.M{"name": "app4"})
}

func CheckUserAccess(teamNames []string, u *User) bool {
	q := bson.M{"_id": bson.M{"$in": teamNames}}
	var teams []Team
	conn, err := db.Conn()
	if err != nil {
		log.Errorf("Failed to connect to the database: %s", err)
		return false
	}
	defer conn.Close()
	conn.Teams().Find(q).All(&teams)
	var wg sync.WaitGroup
	found := make(chan bool, len(teams)+1)
	for _, team := range teams {
		wg.Add(1)
		go func(t Team) {
			if t.ContainsUser(u) {
				found <- true
			}
			wg.Done()
		}(team)
	}
	go func() {
		wg.Wait()
		found <- false
	}()
	return <-found
}

func RemoveServiceInstanceHandler(w http.ResponseWriter, r *http.Request, u *auth.User) error {
	name := r.URL.Query().Get(":name")
	si, err := getServiceInstanceOrError(name, u)
	if err != nil {
		return err
	}
	if len(si.Apps) > 0 {
		msg := "This service instance is bound to at least one app. Unbind them before removing it"
		return &errors.Http{Code: http.StatusInternalServerError, Message: msg}
	}
	if err = si.Service().ProductionEndpoint().Destroy(&si); err != nil {
		return err
	}
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	err = conn.ServiceInstances().Remove(bson.M{"name": name})
	if err != nil {
		return err
	}
	w.Write([]byte("service instance successfuly removed"))
	return nil
}

func ServiceInstanceStatusHandler(w http.ResponseWriter, r *http.Request, u *auth.User) error {
	// #TODO (flaviamissi) should check if user has access to service
	// just call GetServiceInstanceOrError should be enough
	siName := r.URL.Query().Get(":instance")
	var si service.ServiceInstance
	if siName == "" {
		return &errors.Http{Code: http.StatusBadRequest, Message: "Service instance name not provided."}
	}
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	err = conn.ServiceInstances().Find(bson.M{"name": siName}).One(&si)
	if err != nil {
		msg := fmt.Sprintf("Service instance does not exists, error: %s", err.Error())
		return &errors.Http{Code: http.StatusInternalServerError, Message: msg}
	}
	s := si.Service()
	var b string
	if b, err = s.ProductionEndpoint().Status(&si); err != nil {
		msg := fmt.Sprintf("Could not retrieve status of service instance, error: %s", err.Error())
		return &errors.Http{Code: http.StatusInternalServerError, Message: msg}
	}
	b = fmt.Sprintf(`Service instance "%s" is %s`, siName, b)
	n, err := w.Write([]byte(b))
	if n != len(b) {
		return &errors.Http{Code: http.StatusInternalServerError, Message: "Failed to write response body"}
	}
	return nil
}

func ServiceInfoHandler(w http.ResponseWriter, r *http.Request, u *auth.User) error {
	serviceName := r.URL.Query().Get(":name")
	_, err := getServiceOrError(serviceName, u)
	if err != nil {
		return err
	}
	instances := []service.ServiceInstance{}
	teams, err := u.Teams()
	if err != nil {
		return err
	}
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	teamsNames := auth.GetTeamsNames(teams)
	q := bson.M{"service_name": serviceName, "teams": bson.M{"$in": teamsNames}}
	err = conn.ServiceInstances().Find(q).All(&instances)
	if err != nil {
		return err
	}
	b, err := json.Marshal(instances)
	if err != nil {
		return nil
	}
	w.Write(b)
	return nil
}