Golang db.NewSessionProvider函數代碼示例

func TestKerberos(t *testing.T) {
	testutil.VerifyTestType(t, testutil.KerberosTestType)

	Convey("Should be able to run mongoexport with Kerberos auth", t, func() {
		opts, err := testutil.GetKerberosOptions()

		So(err, ShouldBeNil)

		sessionProvider, err := db.NewSessionProvider(*opts)
		So(err, ShouldBeNil)

		export := MongoExport{
			ToolOptions:     *opts,
			OutputOpts:      &OutputFormatOptions{},
			InputOpts:       &InputOptions{},
			SessionProvider: sessionProvider,
		}

		var out bytes.Buffer
		num, err := export.exportInternal(&out)

		So(err, ShouldBeNil)
		So(num, ShouldEqual, 1)
		outputLines := strings.Split(strings.TrimSpace(out.String()), "\n")
		So(len(outputLines), ShouldEqual, 1)
		outMap := map[string]interface{}{}
		So(json.Unmarshal([]byte(outputLines[0]), &outMap), ShouldBeNil)
		So(outMap["kerberos"], ShouldEqual, true)
		So(outMap["authenticated"], ShouldEqual, "yeah")
		So(outMap["_id"].(map[string]interface{})["$oid"], ShouldEqual, "528fb35afb3a8030e2f643c3")
	})
}

// put in some test data into GridFS
func setUpGridFSTestData() ([]interface{}, error) {
	sessionProvider, err := db.NewSessionProvider(*toolOptions)
	if err != nil {
		return nil, err
	}
	session, err := sessionProvider.GetSession()
	if err != nil {
		return nil, err
	}
	defer session.Close()

	bytesExpected := []interface{}{}
	gfs := session.DB(testDB).GridFS("fs")

	var testfile *mgo.GridFile

	for i, item := range []string{"testfile1", "testfile2", "testfile3"} {
		testfile, err = gfs.Create(item)
		if err != nil {
			return nil, err
		}
		defer testfile.Close()

		n, err := testfile.Write([]byte(strings.Repeat("a", (i+1)*5)))
		if err != nil {
			return nil, err
		}

		bytesExpected = append(bytesExpected, n)
	}

	return bytesExpected, nil
}

// Init performs preliminary setup operations for MongoDump.
func (dump *MongoDump) Init() error {
	err := dump.ValidateOptions()
	if err != nil {
		return fmt.Errorf("bad option: %v", err)
	}
	if dump.stdout == nil {
		dump.stdout = os.Stdout
	}
	dump.sessionProvider, err = db.NewSessionProvider(*dump.ToolOptions)
	if err != nil {
		return fmt.Errorf("can't create session: %v", err)
	}

	// temporarily allow secondary reads for the isMongos check
	dump.sessionProvider.SetReadPreference(mgo.Nearest)
	dump.isMongos, err = dump.sessionProvider.IsMongos()
	if err != nil {
		return err
	}

	if dump.isMongos && dump.OutputOptions.Oplog {
		return fmt.Errorf("can't use --oplog option when dumping from a mongos")
	}

	var mode mgo.Mode
	if dump.ToolOptions.ReplicaSetName != "" || dump.isMongos {
		mode = mgo.Primary
	} else {
		mode = mgo.Nearest
	}
	var tags bson.D

	if dump.InputOptions.ReadPreference != "" {
		mode, tags, err = db.ParseReadPreference(dump.InputOptions.ReadPreference)
		if err != nil {
			return fmt.Errorf("error parsing --readPreference : %v", err)
		}
		if len(tags) > 0 {
			dump.sessionProvider.SetTags(tags)
		}
	}

	// warn if we are trying to dump from a secondary in a sharded cluster
	if dump.isMongos && mode != mgo.Primary {
		log.Logf(log.Always, db.WarningNonPrimaryMongosConnection)
	}

	dump.sessionProvider.SetReadPreference(mode)
	dump.sessionProvider.SetTags(tags)
	dump.sessionProvider.SetFlags(db.DisableSocketTimeout)

	// return a helpful error message for mongos --repair
	if dump.OutputOptions.Repair && dump.isMongos {
		return fmt.Errorf("--repair flag cannot be used on a mongos")
	}

	dump.manager = intents.NewIntentManager()
	dump.progressManager = progress.NewProgressBarManager(log.Writer(0), progressBarWaitTime)
	return nil
}

func TestKerberos(t *testing.T) {
	testutil.VerifyTestType(t, testutil.KerberosTestType)

	Convey("Should be able to run mongoexport with Kerberos auth", t, func() {
		opts, err := testutil.GetKerberosOptions()

		So(err, ShouldBeNil)

		sessionProvider, err := db.NewSessionProvider(*opts)
		So(err, ShouldBeNil)

		export := MongoExport{
			ToolOptions:     *opts,
			OutputOpts:      &OutputFormatOptions{},
			InputOpts:       &InputOptions{},
			SessionProvider: sessionProvider,
		}

		var out bytes.Buffer
		num, err := export.exportInternal(&out)

		So(err, ShouldBeNil)
		So(num, ShouldEqual, 1)
		outputLines := strings.Split(strings.TrimSpace(out.String()), "\n")
		So(len(outputLines), ShouldEqual, 1)
		So(outputLines[0], ShouldEqual,
			"{\"_id\":{\"$oid\":\"528fb35afb3a8030e2f643c3\"},"+
				"\"authenticated\":\"yeah\",\"kerberos\":true}")
	})
}

// Init performs preliminary setup operations for MongoDump.
func (dump *MongoDump) Init() error {
	err := dump.ValidateOptions()
	if err != nil {
		return fmt.Errorf("bad option: %v", err)
	}
	dump.sessionProvider, err = db.NewSessionProvider(*dump.ToolOptions)
	if err != nil {
		return fmt.Errorf("can't create session: %v", err)
	}

	// allow secondary reads for the isMongos check
	dump.sessionProvider.SetFlags(db.Monotonic)
	dump.isMongos, err = dump.sessionProvider.IsMongos()
	if err != nil {
		return err
	}

	// ensure we allow secondary reads on mongods and disable TCP timeouts
	flags := db.DisableSocketTimeout
	if dump.isMongos {
		log.Logf(log.Info, "connecting to mongos; secondary reads disabled")
	} else {
		flags |= db.Monotonic
	}
	dump.sessionProvider.SetFlags(flags)

	// return a helpful error message for mongos --repair
	if dump.OutputOptions.Repair && dump.isMongos {
		return fmt.Errorf("--repair flag cannot be used on a mongos")
	}

	dump.manager = intents.NewIntentManager()
	dump.progressManager = progress.NewProgressBarManager(log.Writer(0), progressBarWaitTime)
	return nil
}

func main() {
	go signals.Handle()
	// initialize command-line opts
	opts := options.New("mongofiles", mongofiles.Usage, options.EnabledOptions{Auth: true, Connection: true, Namespace: false})

	storageOpts := &mongofiles.StorageOptions{}
	opts.AddOptions(storageOpts)

	args, err := opts.Parse()
	if err != nil {
		log.Logf(log.Always, "error parsing command line options: %v", err)
		log.Logf(log.Always, "try 'mongofiles --help' for more information")
		os.Exit(util.ExitBadOptions)
	}

	// print help, if specified
	if opts.PrintHelp(false) {
		return
	}

	// print version, if specified
	if opts.PrintVersion() {
		return
	}
	log.SetVerbosity(opts.Verbosity)

	// add the specified database to the namespace options struct
	opts.Namespace.DB = storageOpts.DB

	// connect directly, unless a replica set name is explicitly specified
	_, setName := util.ParseConnectionString(opts.Host)
	opts.Direct = (setName == "")
	opts.ReplicaSetName = setName

	// create a session provider to connect to the db
	provider, err := db.NewSessionProvider(*opts)
	if err != nil {
		log.Logf(log.Always, "error connecting to host: %v", err)
		os.Exit(util.ExitError)
	}
	mf := mongofiles.MongoFiles{
		ToolOptions:     opts,
		StorageOptions:  storageOpts,
		SessionProvider: provider,
	}

	if err := mf.ValidateCommand(args); err != nil {
		log.Logf(log.Always, "%v", err)
		log.Logf(log.Always, "try 'mongofiles --help' for more information")
		os.Exit(util.ExitBadOptions)
	}

	output, err := mf.Run(true)
	if err != nil {
		log.Logf(log.Always, "Failed: %v", err)
		os.Exit(util.ExitError)
	}
	fmt.Printf("%s", output)
}

//Utility constructor for NodeMonitor that copies the same connection settings
//from an instance of ToolOptions, but for a different host name.
func NewNodeMonitor(opts commonopts.ToolOptions, fullHost string) *NodeMonitor {
	optsCopy := opts
	host, port := parseHostPort(fullHost)
	optsCopy.Connection = &commonopts.Connection{host, port}
	optsCopy.Direct = true
	sessionProvider := db.NewSessionProvider(optsCopy)
	return &NodeMonitor{
		host:            fullHost,
		sessionProvider: sessionProvider,
		LastStatus:      nil,
		LastUpdate:      time.Now(),
		Err:             nil,
	}
}

func NewMongoImport() (*MongoImport, error) {
	toolOptions := getBasicToolOptions()
	inputOptions := &InputOptions{}
	ingestOptions := &IngestOptions{}
	provider, err := db.NewSessionProvider(*toolOptions)
	if err != nil {
		return nil, err
	}
	return &MongoImport{
		ToolOptions:     toolOptions,
		InputOptions:    inputOptions,
		IngestOptions:   ingestOptions,
		SessionProvider: provider,
	}, nil
}

func simpleMongoFilesInstance(args []string) (*MongoFiles, error) {
	sessionProvider, err := db.NewSessionProvider(*toolOptions)
	if err != nil {
		return nil, err
	}

	mongofiles := MongoFiles{
		ToolOptions:     toolOptions,
		StorageOptions:  &StorageOptions{GridFSPrefix: "fs", DB: testDB},
		SessionProvider: sessionProvider,
		Command:         args[0],
		FileName:        args[1],
	}

	return &mongofiles, nil
}

// remove test data from GridFS
func tearDownGridFSTestData() error {
	sessionProvider, err := db.NewSessionProvider(*toolOptions)
	if err != nil {
		return err
	}
	session, err := sessionProvider.GetSession()
	if err != nil {
		return err
	}
	defer session.Close()

	if err = session.DB(testDB).DropDatabase(); err != nil {
		return err
	}

	return nil
}

// NewNodeMonitor copies the same connection settings from an instance of
// ToolOptions, but monitors fullHost.
func NewNodeMonitor(opts options.ToolOptions, fullHost string, all bool) (*NodeMonitor, error) {
	optsCopy := opts
	host, port := parseHostPort(fullHost)
	optsCopy.Connection = &options.Connection{Host: host, Port: port}
	optsCopy.Direct = true
	sessionProvider, err := db.NewSessionProvider(optsCopy)
	if err != nil {
		return nil, err
	}
	return &NodeMonitor{
		host:            fullHost,
		sessionProvider: sessionProvider,
		LastStatus:      nil,
		LastUpdate:      time.Now(),
		All:             all,
		Err:             nil,
	}, nil
}

func getBareSession() (*mgo.Session, error) {
	ssl := testutil.GetSSLOptions()
	auth := testutil.GetAuthOptions()
	sessionProvider, err := db.NewSessionProvider(options.ToolOptions{
		Connection: &options.Connection{
			Host: testServer,
			Port: testPort,
		},
		Auth: &auth,
		SSL:  &ssl,
	})
	if err != nil {
		return nil, err
	}
	session, err := sessionProvider.GetSession()
	if err != nil {
		return nil, err
	}
	return session, nil
}

func main() {
	go signals.Handle()

	// initialize command line options
	opts := options.New("mongooplog", mongooplog.Usage,
		options.EnabledOptions{Auth: true, Connection: true, Namespace: false})

	// add the mongooplog-specific options
	sourceOpts := &mongooplog.SourceOptions{}
	opts.AddOptions(sourceOpts)

	log.Logf(log.Always, "warning: mongooplog is deprecated, and will be removed completely in a future release")

	// parse the command line options
	args, err := opts.Parse()
	if err != nil {
		log.Logf(log.Always, "error parsing command line options: %v", err)
		log.Logf(log.Always, "try 'mongooplog --help' for more information")
		os.Exit(util.ExitBadOptions)
	}

	if len(args) != 0 {
		log.Logf(log.Always, "positional arguments not allowed: %v", args)
		log.Logf(log.Always, "try 'mongooplog --help' for more information")
		os.Exit(util.ExitBadOptions)
	}

	// print help, if specified
	if opts.PrintHelp(false) {
		return
	}

	// print version, if specified
	if opts.PrintVersion() {
		return
	}

	// init logger
	log.SetVerbosity(opts.Verbosity)

	// connect directly, unless a replica set name is explicitly specified
	_, setName := util.ParseConnectionString(opts.Host)
	opts.Direct = (setName == "")
	opts.ReplicaSetName = setName

	// validate the mongooplog options
	if sourceOpts.From == "" {
		log.Logf(log.Always, "command line error: need to specify --from")
		os.Exit(util.ExitBadOptions)
	}

	// create a session provider for the destination server
	sessionProviderTo, err := db.NewSessionProvider(*opts)
	if err != nil {
		log.Logf(log.Always, "error connecting to destination host: %v", err)
		os.Exit(util.ExitError)
	}

	// create a session provider for the source server
	opts.Connection.Host = sourceOpts.From
	opts.Connection.Port = ""
	sessionProviderFrom, err := db.NewSessionProvider(*opts)
	if err != nil {
		log.Logf(log.Always, "error connecting to source host: %v", err)
		os.Exit(util.ExitError)
	}

	// initialize mongooplog
	oplog := mongooplog.MongoOplog{
		ToolOptions:         opts,
		SourceOptions:       sourceOpts,
		SessionProviderFrom: sessionProviderFrom,
		SessionProviderTo:   sessionProviderTo,
	}

	// kick it off
	if err := oplog.Run(); err != nil {
		log.Logf(log.Always, "error: %v", err)
		os.Exit(util.ExitError)
	}

}

func main() {
	// initialize command-line opts
	opts := options.New("mongoexport", mongoexport.Usage,
		options.EnabledOptions{Auth: true, Connection: true, Namespace: true})

	outputOpts := &mongoexport.OutputFormatOptions{}
	opts.AddOptions(outputOpts)
	inputOpts := &mongoexport.InputOptions{}
	opts.AddOptions(inputOpts)

	args, err := opts.Parse()
	if err != nil {
		log.Logvf(log.Always, "error parsing command line options: %v", err)
		log.Logvf(log.Always, "try 'mongoexport --help' for more information")
		os.Exit(util.ExitBadOptions)
	}
	if len(args) != 0 {
		log.Logvf(log.Always, "too many positional arguments: %v", args)
		log.Logvf(log.Always, "try 'mongoexport --help' for more information")
		os.Exit(util.ExitBadOptions)
	}

	log.SetVerbosity(opts.Verbosity)
	signals.Handle()

	// print help, if specified
	if opts.PrintHelp(false) {
		return
	}

	// print version, if specified
	if opts.PrintVersion() {
		return
	}

	// connect directly, unless a replica set name is explicitly specified
	_, setName := util.ParseConnectionString(opts.Host)
	opts.Direct = (setName == "")
	opts.ReplicaSetName = setName

	provider, err := db.NewSessionProvider(*opts)

	// temporarily allow secondary reads for the isMongos check
	provider.SetReadPreference(mgo.Nearest)
	isMongos, err := provider.IsMongos()
	if err != nil {
		log.Logvf(log.Always, "%v", err)
		os.Exit(util.ExitError)
	}

	provider.SetFlags(db.DisableSocketTimeout)

	if inputOpts.SlaveOk {
		if inputOpts.ReadPreference != "" {
			log.Logvf(log.Always, "--slaveOk can't be specified when --readPreference is specified")
			os.Exit(util.ExitBadOptions)
		}
		log.Logvf(log.Always, "--slaveOk is deprecated and being internally rewritten as --readPreference=nearest")
		inputOpts.ReadPreference = "nearest"
	}

	var mode mgo.Mode
	if opts.ReplicaSetName != "" || isMongos {
		mode = mgo.Primary
	} else {
		mode = mgo.Nearest
	}
	var tags bson.D
	if inputOpts.ReadPreference != "" {
		mode, tags, err = db.ParseReadPreference(inputOpts.ReadPreference)
		if err != nil {
			log.Logvf(log.Always, "error parsing --ReadPreference: %v", err)
			os.Exit(util.ExitBadOptions)
		}
		if len(tags) > 0 {
			provider.SetTags(tags)
		}
	}

	// warn if we are trying to export from a secondary in a sharded cluster
	if isMongos && mode != mgo.Primary {
		log.Logvf(log.Always, db.WarningNonPrimaryMongosConnection)
	}

	provider.SetReadPreference(mode)

	if err != nil {
		log.Logvf(log.Always, "error connecting to host: %v", err)
		os.Exit(util.ExitError)
	}
	exporter := mongoexport.MongoExport{
		ToolOptions:     *opts,
		OutputOpts:      outputOpts,
		InputOpts:       inputOpts,
		SessionProvider: provider,
	}

	err = exporter.ValidateSettings()
	if err != nil {
		log.Logvf(log.Always, "error validating settings: %v", err)
//.........這裏部分代碼省略.........

func TestCollectionExists(t *testing.T) {

	testutil.VerifyTestType(t, testutil.IntegrationTestType)

	Convey("With a test mongorestore", t, func() {
		ssl := testutil.GetSSLOptions()
		auth := testutil.GetAuthOptions()
		sessionProvider, err := db.NewSessionProvider(commonOpts.ToolOptions{
			Connection: &commonOpts.Connection{
				Host: "localhost",
				Port: db.DefaultTestPort,
			},
			Auth: &auth,
			SSL:  &ssl,
		})
		So(err, ShouldBeNil)

		restore := &MongoRestore{
			SessionProvider: sessionProvider,
		}

		Convey("and some test data in a server", func() {
			session, err := restore.SessionProvider.GetSession()
			So(err, ShouldBeNil)
			So(session.DB(ExistsDB).C("one").Insert(bson.M{}), ShouldBeNil)
			So(session.DB(ExistsDB).C("two").Insert(bson.M{}), ShouldBeNil)
			So(session.DB(ExistsDB).C("three").Insert(bson.M{}), ShouldBeNil)

			Convey("collections that exist should return true", func() {
				exists, err := restore.CollectionExists(&intents.Intent{DB: ExistsDB, C: "one"})
				So(err, ShouldBeNil)
				So(exists, ShouldBeTrue)
				exists, err = restore.CollectionExists(&intents.Intent{DB: ExistsDB, C: "two"})
				So(err, ShouldBeNil)
				So(exists, ShouldBeTrue)
				exists, err = restore.CollectionExists(&intents.Intent{DB: ExistsDB, C: "three"})
				So(err, ShouldBeNil)
				So(exists, ShouldBeTrue)

				Convey("and those that do not exist should return false", func() {
					exists, err = restore.CollectionExists(&intents.Intent{DB: ExistsDB, C: "four"})
					So(err, ShouldBeNil)
					So(exists, ShouldBeFalse)
				})
			})

			Reset(func() {
				session.DB(ExistsDB).DropDatabase()
				session.Close()
			})
		})

		Convey("and a fake cache should be used instead of the server when it exists", func() {
			restore.knownCollections = map[string][]string{
				ExistsDB: []string{"cats", "dogs", "snakes"},
			}
			exists, err := restore.CollectionExists(&intents.Intent{DB: ExistsDB, C: "dogs"})
			So(err, ShouldBeNil)
			So(exists, ShouldBeTrue)
			exists, err = restore.CollectionExists(&intents.Intent{DB: ExistsDB, C: "two"})
			So(err, ShouldBeNil)
			So(exists, ShouldBeFalse)
		})
	})
}