Skip to content

Commit f22b29f

Browse files
sushanbnimf
sushanb
and
nimf
authored
feat(bigtable): expose a client option to disable direct access: (#14626)
Co-authored-by: Yuri Golobokov <yuri.golobokov@icloud.com>
1 parent b291ee8 commit f22b29f

3 files changed

Lines changed: 207 additions & 51 deletions

File tree

‎bigtable/client.go‎

Lines changed: 65 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,8 @@ import (
3535
"google.golang.org/grpc/metadata"
3636
)
3737

38+
const directpathEnvVar = "CBT_ENABLE_DIRECTPATH"
39+
3840
// Client is a client for reading and writing data to tables in an instance.
3941
//
4042
// A Client is safe to use concurrently, except for its Close method.
@@ -72,6 +74,9 @@ type ClientConfig struct {
7274
// DisableConnectionRecycler disables the automatic preemptive refresh of connection.
7375
// Preemptive connection is default to true
7476
DisableConnectionRecycler bool
77+
78+
// DisableDirectAccess disables direct access by default.
79+
DisableDirectAccess bool
7580
}
7681

7782
// MetricsProvider is a wrapper for built in metrics meter provider
@@ -129,6 +134,7 @@ func NewClientWithConfig(ctx context.Context, project, instance string, config C
129134
var directPathOptions = []option.ClientOption{
130135
internaloption.EnableDirectPath(true),
131136
internaloption.EnableDirectPathXds(),
137+
internaloption.AllowHardBoundTokens("ALTS"),
132138
}
133139

134140
// Allow non-default service account in DirectPath.
@@ -155,10 +161,11 @@ func NewClientWithConfig(ctx context.Context, project, instance string, config C
155161
// as CFE/GFE will call RLS with gslb target type
156162
// only TD calls the RLS with grpc target type
157163
// and we evaluate the directAccess option after that.
158-
directAccessMD := createFeatureFlagsMD(metricsTracerFactory.enabled, disableRetryInfo, true)
164+
165+
allowDirectAccess := isDirectAccessEnabled(config)
166+
directAccessMD := createFeatureFlagsMD(metricsTracerFactory.enabled, disableRetryInfo, allowDirectAccess)
159167

160168
var connPool gtransport.ConnPool
161-
var connPoolErr error
162169
var dsm *btransport.DynamicScaleMonitor
163170
var connRecycler *btransport.ConnectionRecycler
164171

@@ -171,7 +178,18 @@ func NewClientWithConfig(ctx context.Context, project, instance string, config C
171178
}
172179
}
173180
var connPoolSize int
174-
if enableBigtableConnPool {
181+
if !enableBigtableConnPool {
182+
// Use the regular ConnPool
183+
// For regular ConnPool the Direct Access is off by default so we need to check the env var again.
184+
if enabled, _ := strconv.ParseBool(os.Getenv(directpathEnvVar)); enabled {
185+
o = append(o, directPathOptions...)
186+
}
187+
regConnPool, err := gtransport.DialPool(ctx, o...)
188+
if err != nil {
189+
return nil, err
190+
}
191+
connPool = regConnPool
192+
} else { // Use the BigtableConnPool
175193
uResolver, err := internaloption.NewUnsafeResolver(o...)
176194
if err != nil {
177195
// just fallback
@@ -186,18 +204,29 @@ func NewClientWithConfig(ctx context.Context, project, instance string, config C
186204

187205
fullInstanceName := fmt.Sprintf("projects/%s/instances/%s", project, instance)
188206

189-
directAccessDialerOptions := make([]option.ClientOption, len(o))
190-
copy(directAccessDialerOptions, o)
191-
directAccessDialerOptions = append(directAccessDialerOptions, directPathOptions...)
192-
// enable hard bound tokens by default
193-
directAccessDialerOptions = append(directAccessDialerOptions, internaloption.AllowHardBoundTokens("ALTS"))
207+
var poolOpts []btransport.BigtableChannelPoolOption
208+
poolOpts = append(poolOpts,
209+
btransport.WithInstanceName(fullInstanceName),
210+
btransport.WithAppProfile(config.AppProfile),
211+
btransport.WithFeatureFlagsMetadata(directAccessMD),
212+
btransport.WithMetricsReporterConfig(btopt.DefaultMetricsReporterConfig()),
213+
btransport.WithMeterProvider(metricsTracerFactory.otelMeterProvider),
214+
btransport.WithDirectAccessFeatureFlagsMetadata(directAccessMD),
215+
)
194216

195-
directAccessDialer := func() (*btransport.BigtableConn, error) {
196-
grpcConn, err := gtransport.Dial(ctx, directAccessDialerOptions...)
197-
if err != nil {
198-
return nil, err
217+
// Only setup DirectPath dialers if not disabled by config/env
218+
if allowDirectAccess {
219+
directAccessDialerOptions := make([]option.ClientOption, len(o))
220+
copy(directAccessDialerOptions, o)
221+
directAccessDialerOptions = append(directAccessDialerOptions, directPathOptions...)
222+
directAccessDialer := func() (*btransport.BigtableConn, error) {
223+
grpcConn, err := gtransport.Dial(ctx, directAccessDialerOptions...)
224+
if err != nil {
225+
return nil, err
226+
}
227+
return btransport.NewBigtableConn(grpcConn), nil
199228
}
200-
return btransport.NewBigtableConn(grpcConn), nil
229+
poolOpts = append(poolOpts, btransport.WithDirectAccessDialer(directAccessDialer))
201230
}
202231

203232
btPool, err := btransport.NewBigtableChannelPool(ctx,
@@ -212,51 +241,28 @@ func NewClientWithConfig(ctx context.Context, project, instance string, config C
212241
},
213242
clientCreationTimestamp,
214243
// options
215-
btransport.WithInstanceName(fullInstanceName),
216-
btransport.WithAppProfile(config.AppProfile),
217-
btransport.WithFeatureFlagsMetadata(directAccessMD),
218-
btransport.WithMetricsReporterConfig(btopt.DefaultMetricsReporterConfig()),
219-
btransport.WithMeterProvider(metricsTracerFactory.otelMeterProvider),
220-
btransport.WithDirectAccessFeatureFlagsMetadata(directAccessMD),
221-
btransport.WithDirectAccessDialer(directAccessDialer),
244+
poolOpts...,
222245
)
223-
224246
if err != nil {
225-
connPoolErr = err
226-
} else {
227-
connPool = btPool
228-
229-
// Validate dynamic config early if enabled
230-
if !config.DisableDynamicChannelPool {
231-
if err := btransport.ValidateDynamicConfig(btopt.DefaultDynamicChannelPoolConfig(), defaultBigtableConnPoolSize); err != nil {
232-
return nil, fmt.Errorf("invalid DynamicChannelPoolConfig: %w", err)
233-
}
247+
return nil, err
248+
}
234249

235-
dsm = btransport.NewDynamicScaleMonitor(btopt.DefaultDynamicChannelPoolConfig(), btPool)
236-
dsm.Start(ctx) // Start the monitor's background goroutine
237-
}
238-
// connection recyler.
239-
if !config.DisableConnectionRecycler {
240-
connRecycler = btransport.NewConnectionRecycler(btopt.DefaultConnectionRecycleConfig(), btPool)
241-
connRecycler.Start(ctx) // Start the monitor's background goroutine
250+
connPool = btPool
251+
252+
// Validate dynamic config early if enabled
253+
if !config.DisableDynamicChannelPool {
254+
if err := btransport.ValidateDynamicConfig(btopt.DefaultDynamicChannelPoolConfig(), defaultBigtableConnPoolSize); err != nil {
255+
return nil, fmt.Errorf("invalid DynamicChannelPoolConfig: %w", err)
242256
}
243257

258+
dsm = btransport.NewDynamicScaleMonitor(btopt.DefaultDynamicChannelPoolConfig(), btPool)
259+
dsm.Start(ctx) // Start the monitor's background goroutine
244260
}
245-
246-
} else {
247-
enableDirectAccess, _ := strconv.ParseBool(os.Getenv("CBT_ENABLE_DIRECTPATH"))
248-
if enableDirectAccess {
249-
o = append(o, directPathOptions...)
250-
if disableBoundToken, _ := strconv.ParseBool(os.Getenv("CBT_DISABLE_DIRECTPATH_BOUND_TOKEN")); !disableBoundToken {
251-
o = append(o, internaloption.AllowHardBoundTokens("ALTS"))
252-
}
261+
// connection recyler.
262+
if !config.DisableConnectionRecycler {
263+
connRecycler = btransport.NewConnectionRecycler(btopt.DefaultConnectionRecycleConfig(), btPool)
264+
connRecycler.Start(ctx) // Start the monitor's background goroutine
253265
}
254-
// use to regular ConnPool
255-
connPool, connPoolErr = gtransport.DialPool(ctx, o...)
256-
}
257-
258-
if connPoolErr != nil {
259-
return nil, connPoolErr
260266
}
261267

262268
return &Client{
@@ -400,3 +406,11 @@ func (c *Client) newBuiltinMetricsTracer(ctx context.Context, table string, isSt
400406
mt := c.metricsTracerFactory.createBuiltinMetricsTracer(ctx, table, isStreaming)
401407
return &mt
402408
}
409+
410+
func isDirectAccessEnabled(config ClientConfig) bool {
411+
if os.Getenv(directpathEnvVar) == "" {
412+
return !config.DisableDirectAccess
413+
}
414+
res, _ := strconv.ParseBool(os.Getenv(directpathEnvVar))
415+
return res
416+
}

‎bigtable/client_test.go‎

Lines changed: 139 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package bigtable
16+
17+
import (
18+
"os"
19+
"testing"
20+
)
21+
22+
func TestIsDirectAccessDisabled(t *testing.T) {
23+
tests := []struct {
24+
name string
25+
configDisable bool
26+
envValue string
27+
envSet bool
28+
want bool // direct access should be enabled
29+
}{
30+
{
31+
name: "config_disabled_env_unset",
32+
configDisable: true,
33+
envSet: false,
34+
want: false,
35+
},
36+
{
37+
name: "config_disabled_env_false",
38+
configDisable: true,
39+
envSet: true,
40+
envValue: "false",
41+
want: false,
42+
},
43+
{
44+
name: "config_disabled_env_true",
45+
configDisable: true,
46+
envSet: true,
47+
envValue: "true",
48+
want: true,
49+
},
50+
{
51+
name: "config_enabled_env_unset",
52+
configDisable: false,
53+
envSet: false,
54+
want: true,
55+
},
56+
{
57+
name: "config_enabled_env_false_lowercase",
58+
configDisable: false,
59+
envSet: true,
60+
envValue: "false",
61+
want: false,
62+
},
63+
{
64+
name: "config_enabled_env_false_uppercase",
65+
configDisable: false,
66+
envSet: true,
67+
envValue: "FALSE",
68+
want: false,
69+
},
70+
{
71+
name: "config_enabled_env_false_mixedcase",
72+
configDisable: false,
73+
envSet: true,
74+
envValue: "False",
75+
want: false,
76+
},
77+
{
78+
name: "config_enabled_env_true_lowercase",
79+
configDisable: false,
80+
envSet: true,
81+
envValue: "true",
82+
want: true,
83+
},
84+
{
85+
name: "config_enabled_env_true_uppercase",
86+
configDisable: false,
87+
envSet: true,
88+
envValue: "TRUE",
89+
want: true,
90+
},
91+
{
92+
name: "config_enabled_env_true_mixedcase",
93+
configDisable: false,
94+
envSet: true,
95+
envValue: "True",
96+
want: true,
97+
},
98+
{
99+
// 't' is not respected, defaults to not disabled (false)
100+
name: "config_enabled_env_t",
101+
configDisable: false,
102+
envSet: true,
103+
envValue: "t",
104+
want: true,
105+
},
106+
{
107+
// 'f' is not respected, defaults to not disabled (false)
108+
name: "config_enabled_env_f",
109+
configDisable: false,
110+
envSet: true,
111+
envValue: "f",
112+
want: false,
113+
},
114+
{
115+
name: "config_enabled_env_invalid",
116+
configDisable: false,
117+
envSet: true,
118+
envValue: "invalid",
119+
want: false,
120+
},
121+
}
122+
for _, tc := range tests {
123+
t.Run(tc.name, func(t *testing.T) {
124+
if tc.envSet {
125+
os.Setenv("CBT_ENABLE_DIRECTPATH", tc.envValue)
126+
defer os.Unsetenv("CBT_ENABLE_DIRECTPATH")
127+
} else {
128+
os.Unsetenv("CBT_ENABLE_DIRECTPATH")
129+
}
130+
config := ClientConfig{
131+
DisableDirectAccess: tc.configDisable,
132+
}
133+
got := isDirectAccessEnabled(config)
134+
if got != tc.want {
135+
t.Errorf("isDirectAccessDisabled(%+v) = %v; want %v", config, got, tc.want)
136+
}
137+
})
138+
}
139+
}

‎bigtable/internal/transport/connpool.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -479,6 +479,9 @@ func NewBigtableChannelPool(ctx context.Context, connPoolSize int, strategy btop
479479
btopt.Debugf(pool.logger, "bigtable_connpool: Direct Access is not available. Using standard path")
480480
}
481481
}
482+
} else {
483+
btopt.Debugf(pool.logger, "bigtable_connpool: Direct Access manually disabled via config or environment.")
484+
pool.reportDirectAccessFailure("manually_disabled")
482485
}
483486

484487
// Initialize the connectionFactory

0 commit comments

Comments
 (0)