GODRIVER-1507 Correctly pass URI to topology (#320)

Author: Divjot Arora

mongo/client.go

@@ -325,6 +325,11 @@ func (c *Client) configure(opts *options.ClientOptions) error {
 
 	// TODO(GODRIVER-814): Add tests for topology, server, and connection related options.
 
+	// Pass down URI so topology can determine whether or not SRV polling is required
+	topologyOpts = append(topologyOpts, topology.WithURI(func(uri string) string {
+		return opts.GetURI()
+	}))
+
 	// AppName
 	var appName string
 	if opts.AppName != nil {
@@ -559,7 +564,9 @@ func (c *Client) configure(opts *options.ClientOptions) error {
 
 	// Deployment
 	if opts.Deployment != nil {
-		if len(serverOpts) > 2 || len(topologyOpts) > 1 {
+		// topology options: WithSeedlist and WithURI
+		// server options: WithClock and WithConnectionOptions
+		if len(serverOpts) > 2 || len(topologyOpts) > 2 {
 			return errors.New("cannot specify topology or server options with a deployment")
 		}
 		c.deployment = opts.Deployment

mongo/client_test.go

@@ -239,4 +239,22 @@ func TestClient(t *testing.T) {
 		client := setupClient(options.Client().SetWriteConcern(wc))
 		assert.Equal(t, wc, client.writeConcern, "mismatch; expected write concern %v, got %v", wc, client.writeConcern)
 	})
+	t.Run("GetURI", func(t *testing.T) {
+		t.Run("ApplyURI not called", func(t *testing.T) {
+			opts := options.Client().SetHosts([]string{"localhost:27017"})
+			uri := opts.GetURI()
+			assert.Equal(t, "", uri, "expected GetURI to return empty string, got %v", uri)
+		})
+		t.Run("ApplyURI called with empty string", func(t *testing.T) {
+			opts := options.Client().ApplyURI("")
+			uri := opts.GetURI()
+			assert.Equal(t, "", uri, "expected GetURI to return empty string, got %v", uri)
+		})
+		t.Run("ApplyURI called with non-empty string", func(t *testing.T) {
+			uri := "mongodb://localhost:27017/foobar"
+			opts := options.Client().ApplyURI(uri)
+			got := opts.GetURI()
+			assert.Equal(t, uri, got, "expected GetURI to return %v, got %v", uri, got)
+		})
+	})
 }

mongo/options/clientoptions.go

@@ -118,6 +118,7 @@ type ClientOptions struct {
 	AutoEncryptionOptions  *AutoEncryptionOptions
 
 	err error
+	uri string
 
 	// These options are for internal use only and should not be set. They are deprecated and are
 	// not part of the stability guarantee. They may be removed in the future.
@@ -133,6 +134,12 @@ func Client() *ClientOptions {
 // Validate validates the client options. This method will return the first error found.
 func (c *ClientOptions) Validate() error { return c.err }
 
+// GetURI returns the original URI used to configure the ClientOptions instance. If ApplyURI was not called during
+// construction, this returns "".
+func (c *ClientOptions) GetURI() string {
+	return c.uri
+}
+
 // ApplyURI parses the given URI and sets options accordingly. The URI can contain host names, IPv4/IPv6 literals, or
 // an SRV record that will be resolved when the Client is created. When using an SRV record, TLS support is
 // implictly enabled. Specify the "tls=false" URI option to override this.
@@ -152,6 +159,7 @@ func (c *ClientOptions) ApplyURI(uri string) *ClientOptions {
 		return c
 	}
 
+	c.uri = uri
 	cs, err := connstring.Parse(uri)
 	if err != nil {
 		c.err = err

mongo/options/clientoptions_test.go

@@ -439,6 +439,7 @@ func TestClientOptions(t *testing.T) {
 		for _, tc := range testCases {
 			t.Run(tc.name, func(t *testing.T) {
 				result := Client().ApplyURI(tc.uri)
+				tc.result.uri = tc.uri // manually add URI to avoid writing it in each test
 				if diff := cmp.Diff(
 					tc.result, result,
 					cmp.AllowUnexported(readconcern.ReadConcern{}, writeconcern.WriteConcern{}, readpref.ReadPref{}),

x/mongo/driver/topology/polling_srv_records_test.go

@@ -131,7 +131,10 @@ func TestPollingSRVRecordsSpec(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			cs, err := connstring.Parse("mongodb+srv://test1.test.build.10gen.cc/?heartbeatFrequencyMS=100")
 			require.NoError(t, err, "Problem parsing the uri: %v", err)
-			topo, err := New(WithConnString(func(connstring.ConnString) connstring.ConnString { return cs }))
+			topo, err := New(
+				WithConnString(func(connstring.ConnString) connstring.ConnString { return cs }),
+				WithURI(func(string) string { return cs.Original }),
+			)
 			require.NoError(t, err, "Could not create the topology: %v", err)
 			mockRes := newMockResolver(tt.recordsToAdd, tt.recordsToRemove, tt.lookupFail, tt.lookupTimeout)
 			topo.dnsResolver = &dns.Resolver{mockRes.LookupSRV, mockRes.LookupTXT}
@@ -167,7 +170,10 @@ func TestPollSRVRecords(t *testing.T) {
 	t.Run("Not unknown or sharded topology", func(t *testing.T) {
 		cs, err := connstring.Parse("mongodb+srv://test1.test.build.10gen.cc/?heartbeatFrequencyMS=100")
 		require.NoError(t, err, "Problem parsing the uri: %v", err)
-		topo, err := New(WithConnString(func(connstring.ConnString) connstring.ConnString { return cs }))
+		topo, err := New(
+			WithConnString(func(connstring.ConnString) connstring.ConnString { return cs }),
+			WithURI(func(string) string { return cs.Original }),
+		)
 		require.NoError(t, err, "Could not create the topology: %v", err)
 		mockRes := newMockResolver(nil, nil, false, false)
 		topo.dnsResolver = &dns.Resolver{mockRes.LookupSRV, mockRes.LookupTXT}
@@ -209,7 +215,10 @@ func TestPollSRVRecords(t *testing.T) {
 	t.Run("Failed Hostname Verification", func(t *testing.T) {
 		cs, err := connstring.Parse("mongodb+srv://test1.test.build.10gen.cc/?heartbeatFrequencyMS=100")
 		require.NoError(t, err, "Problem parsing the uri: %v", err)
-		topo, err := New(WithConnString(func(connstring.ConnString) connstring.ConnString { return cs }))
+		topo, err := New(
+			WithConnString(func(connstring.ConnString) connstring.ConnString { return cs }),
+			WithURI(func(string) string { return cs.Original }),
+		)
 		require.NoError(t, err, "Could not create the topology: %v", err)
 		mockRes := newMockResolver([]*net.SRV{{"blah.bleh", 27019, 0, 0}, {"localhost.test.build.10gen.cc.", 27020, 0, 0}}, nil, false, false)
 		topo.dnsResolver = &dns.Resolver{mockRes.LookupSRV, mockRes.LookupTXT}
@@ -240,7 +249,10 @@ func TestPollSRVRecords(t *testing.T) {
 	t.Run("Return to polling time", func(t *testing.T) {
 		cs, err := connstring.Parse("mongodb+srv://test1.test.build.10gen.cc/?heartbeatFrequencyMS=100")
 		require.NoError(t, err, "Problem parsing the uri: %v", err)
-		topo, err := New(WithConnString(func(connstring.ConnString) connstring.ConnString { return cs }))
+		topo, err := New(
+			WithConnString(func(connstring.ConnString) connstring.ConnString { return cs }),
+			WithURI(func(string) string { return cs.Original }),
+		)
 		require.NoError(t, err, "Could not create the topology: %v", err)
 		mockRes := newMockResolver(nil, nil, false, false)
 		mockRes.fail = 1

x/mongo/driver/topology/topology.go

@@ -64,6 +64,7 @@ type Topology struct {
 
 	done chan struct{}
 
+	pollingRequired   bool
 	pollingDone       chan struct{}
 	pollingwg         sync.WaitGroup
 	rescanSRVInterval time.Duration
@@ -131,6 +132,10 @@ func New(opts ...Option) (*Topology, error) {
 		t.fsm.Kind = description.Single
 	}
 
+	if t.cfg.uri != "" {
+		t.pollingRequired = strings.HasPrefix(t.cfg.uri, "mongodb+srv://")
+	}
+
 	return t, nil
 }
 
@@ -154,7 +159,7 @@ func (t *Topology) Connect() error {
 	}
 	t.serversLock.Unlock()
 
-	if srvPollingRequired(t.cfg.cs.Original) {
+	if t.pollingRequired {
 		go t.pollSRVRecords()
 		t.pollingwg.Add(1)
 	}
@@ -192,7 +197,7 @@ func (t *Topology) Disconnect(ctx context.Context) error {
 	t.subscriptionsClosed = true
 	t.subLock.Unlock()
 
-	if srvPollingRequired(t.cfg.cs.Original) {
+	if t.pollingRequired {
 		t.pollingDone <- struct{}{}
 		t.pollingwg.Wait()
 	}
@@ -203,10 +208,6 @@ func (t *Topology) Disconnect(ctx context.Context) error {
 	return nil
 }
 
-func srvPollingRequired(connstr string) bool {
-	return strings.HasPrefix(connstr, "mongodb+srv://")
-}
-
 // Description returns a description of the topology.
 func (t *Topology) Description() description.Topology {
 	td, ok := t.desc.Load().(description.Topology)
@@ -498,7 +499,7 @@ func (t *Topology) pollSRVRecords() {
 	}()
 
 	// remove the scheme
-	uri := t.cfg.cs.Original[14:]
+	uri := t.cfg.uri[14:]
 	hosts := uri
 	if idx := strings.IndexAny(uri, "/?@"); idx != -1 {
 		hosts = uri[:idx]

x/mongo/driver/topology/topology_options.go

@@ -31,7 +31,8 @@ type config struct {
 	replicaSetName         string
 	seedList               []string
 	serverOpts             []ServerOption
-	cs                     connstring.ConnString
+	cs                     connstring.ConnString // This must not be used for any logic in topology.Topology.
+	uri                    string
 	serverSelectionTimeout time.Duration
 }
 
@@ -275,6 +276,14 @@ func WithServerSelectionTimeout(fn func(time.Duration) time.Duration) Option {
 	}
 }
 
+// WithURI specifies the URI that was used to create the topology.
+func WithURI(fn func(string) string) Option {
+	return func(cfg *config) error {
+		cfg.uri = fn(cfg.uri)
+		return nil
+	}
+}
+
 // addCACertFromFile adds a root CA certificate to the configuration given a path
 // to the containing file.
 func addCACertFromFile(cfg *tls.Config, file string) error {

x/mongo/driver/topology/topology_test.go

@@ -514,3 +514,28 @@ func TestTopology_String_Race(t *testing.T) {
 	<-ch
 	<-ch
 }
+
+func TestTopologyConstruction(t *testing.T) {
+	t.Run("construct with URI", func(t *testing.T) {
+		testCases := []struct {
+			name            string
+			uri             string
+			pollingRequired bool
+		}{
+			{"normal", "mongodb://localhost:27017", false},
+			{"srv", "mongodb+srv://localhost:27017", true},
+		}
+		for _, tc := range testCases {
+			t.Run(tc.name, func(t *testing.T) {
+				topo, err := New(
+					WithURI(func(string) string { return tc.uri }),
+				)
+				assert.Nil(t, err, "topology.New error: %v", err)
+
+				assert.Equal(t, tc.uri, topo.cfg.uri, "expected topology URI to be %v, got %v", tc.uri, topo.cfg.uri)
+				assert.Equal(t, tc.pollingRequired, topo.pollingRequired,
+					"expected topo.pollingRequired to be %v, got %v", tc.pollingRequired, topo.pollingRequired)
+			})
+		}
+	})
+}