Fix database user auto-provisioning

e2e/aws/databases_test.go

@@ -129,12 +129,8 @@ const (
 func postgresConnTest(t *testing.T, cluster *helpers.TeleInstance, user string, route tlsca.RouteToDatabase, query string) {
 	t.Helper()
 	var pgConn *pgconn.PgConn
-	// retry for a while, the database service might need time to give
-	// itself IAM rds:connect permissions.
-	require.EventuallyWithT(t, func(t *assert.CollectT) {
+	waitForDBConnection(t, func(ctx context.Context) error {
 		var err error
-		ctx, cancel := context.WithTimeout(context.Background(), connRetryTick)
-		defer cancel()
 		pgConn, err = postgres.MakeTestClient(ctx, common.TestClientConfig{
 			AuthClient:      cluster.GetSiteAPI(cluster.Secrets.SiteName),
 			AuthServer:      cluster.Process.GetAuthServer(),
@@ -143,10 +139,8 @@ func postgresConnTest(t *testing.T, cluster *helpers.TeleInstance, user string,
 			Username:        user,
 			RouteToDatabase: route,
 		})
-		assert.NoError(t, err)
-		assert.NotNil(t, pgConn)
-	}, waitForConnTimeout, connRetryTick, "connecting to postgres")
-
+		return err
+	})
 	execPGTestQuery(t, pgConn, query)
 }
 
@@ -161,17 +155,11 @@ func postgresLocalProxyConnTest(t *testing.T, cluster *helpers.TeleInstance, use
 	pgconnConfig.User = route.Username
 	pgconnConfig.Database = route.Database
 	var pgConn *pgconn.PgConn
-	// retry for a while, the database service might need time to give
-	// itself IAM rds:connect permissions.
-	require.EventuallyWithT(t, func(t *assert.CollectT) {
+	waitForDBConnection(t, func(ctx context.Context) error {
 		var err error
-		ctx, cancel := context.WithTimeout(context.Background(), connRetryTick)
-		defer cancel()
 		pgConn, err = pgconn.ConnectConfig(ctx, pgconnConfig)
-		assert.NoError(t, err)
-		assert.NotNil(t, pgConn)
-	}, waitForConnTimeout, connRetryTick, "connecting to postgres")
-
+		return err
+	})
 	execPGTestQuery(t, pgConn, query)
 }
 
@@ -204,23 +192,18 @@ func mysqlLocalProxyConnTest(t *testing.T, cluster *helpers.TeleInstance, user s
 	lp := startLocalALPNProxy(t, user, cluster, route)
 
 	var conn *mysqlclient.Conn
-	// retry for a while, the database service might need time to give
-	// itself IAM rds:connect permissions.
-	require.EventuallyWithT(t, func(t *assert.CollectT) {
+	waitForDBConnection(t, func(ctx context.Context) error {
 		var err error
 		var nd net.Dialer
-		ctx, cancel := context.WithTimeout(context.Background(), connRetryTick)
-		defer cancel()
 		conn, err = mysqlclient.ConnectWithDialer(ctx, "tcp",
 			lp.GetAddr(),
 			route.Username,
 			"", /*no password*/
 			route.Database,
 			nd.DialContext,
 		)
-		assert.NoError(t, err)
-		assert.NotNil(t, conn)
-	}, waitForConnTimeout, connRetryTick, "connecting to mysql")
+		return err
+	})
 	defer func() {
 		// Disconnect.
 		require.NoError(t, conn.Close())
@@ -413,6 +396,10 @@ func (c *pgConn) Exec(ctx context.Context, sql string, args ...interface{}) (pgc
 		out, err = c.Conn.Exec(ctx, sql, args...)
 		return trace.Wrap(err)
 	})
+	c.logger.InfoContext(ctx, "Executed sql statement",
+		"sql", sql,
+		"error", err,
+	)
 	return out, trace.Wrap(err)
 }
 
@@ -473,3 +460,28 @@ func isRetryable(err error) bool {
 	}
 	return pgconn.SafeToRetry(err)
 }
+
+func waitForDBConnection(t *testing.T, connectFn func(context.Context) error) {
+	t.Helper()
+	// retry for a while, the database service might need time to give itself
+	// IAM permissions.
+	waitForSuccess(t, func() error {
+		ctx, cancel := context.WithTimeout(context.Background(), connRetryTick)
+		defer cancel()
+		return connectFn(ctx)
+	}, waitForConnTimeout, connRetryTick, "connecting to database")
+}
+
+// waitForSuccess is a test helper that wraps require.EventuallyWithT but runs
+// the given fn first to avoid waiting for the first timer tick.
+func waitForSuccess(t *testing.T, fn func() error, waitDur, tick time.Duration, msgAndArgs ...any) {
+	t.Helper()
+	// EventuallyWithT waits for the first tick before it makes the first
+	// attempt, so to speed things up we check for fn success first.
+	if err := fn(); err == nil {
+		return
+	}
+	require.EventuallyWithT(t, func(t *assert.CollectT) {
+		assert.NoError(t, fn())
+	}, waitDur, tick, msgAndArgs...)
+}

e2e/aws/fixtures_test.go

@@ -159,7 +159,7 @@ func createTeleportCluster(t *testing.T, opts ...testOptionsFunc) *helpers.TeleI
 		teleport.AddUserWithRole(name, roles...)
 	}
 
-	tconf := newTeleportConfig(t)
+	tconf := newTeleportConfig()
 	for _, optFn := range options.serviceConfigFuncs {
 		optFn(tconf)
 	}
@@ -194,7 +194,7 @@ func newInstanceConfig(t *testing.T) helpers.InstanceConfig {
 	}
 }
 
-func newTeleportConfig(t *testing.T) *servicecfg.Config {
+func newTeleportConfig() *servicecfg.Config {
 	tconf := servicecfg.MakeDefaultConfig()
 	// Replace the default auth and proxy listeners with the ones so we can
 	// run multiple tests in parallel.
@@ -207,13 +207,13 @@ func newTeleportConfig(t *testing.T) *servicecfg.Config {
 
 // withUserRole creates a new role that will be bootstraped and then granted to
 // the Teleport user under test.
-func withUserRole(t *testing.T, user, name string, spec types.RoleSpecV6) testOptionsFunc {
+func withUserRole(t *testing.T, userName, roleName string, spec types.RoleSpecV6) testOptionsFunc {
 	t.Helper()
 	// Create a new role with full access to all databases.
-	role, err := types.NewRole(name, spec)
+	role, err := types.NewRole(roleName, spec)
 	require.NoError(t, err)
 	return func(options *testOptions) {
-		options.userRoles[user] = append(options.userRoles[user], role)
+		options.userRoles[userName] = append(options.userRoles[userName], role)
 	}
 }
 
@@ -291,3 +291,16 @@ func makeAutoUserDropRoleSpec(roles ...string) types.RoleSpecV6 {
 	spec.Options.CreateDatabaseUserMode = types.CreateDatabaseUserMode_DB_USER_MODE_BEST_EFFORT_DROP
 	return spec
 }
+
+func makeAutoUserDBPermissions(dbPermissions ...types.DatabasePermission) types.RoleSpecV6 {
+	return types.RoleSpecV6{
+		Allow: types.RoleConditions{
+			DatabaseLabels:      types.Labels{types.Wildcard: []string{types.Wildcard}},
+			DatabaseNames:       []string{types.Wildcard},
+			DatabasePermissions: dbPermissions,
+		},
+		Options: types.RoleOptions{
+			CreateDatabaseUserMode: types.CreateDatabaseUserMode_DB_USER_MODE_KEEP,
+		},
+	}
+}

e2e/aws/main_test.go

@@ -23,6 +23,7 @@ import (
 
 	"github.com/gravitational/teleport/integration/helpers"
 	"github.com/gravitational/teleport/lib"
+	"github.com/gravitational/teleport/lib/utils"
 )
 
 const (
@@ -140,6 +141,7 @@ const (
 // TestMain will re-execute Teleport to run a command if "exec" is passed to
 // it as an argument. Otherwise, it will run tests as normal.
 func TestMain(m *testing.M) {
+	utils.InitLoggerForTests()
 	// agents connect over a reverse tunnel to proxy, so we use insecure mode.
 	lib.SetInsecureDevMode(true)
 	helpers.TestMainImplementation(m)