Skip to content

Commit 26d5493

Browse files
colin-rogers-dbtcolin-rogers-dbt
authored
add support for tcp_keepalive (#1194)
1 parent 4bdcaf6 commit 26d5493

File tree

7 files changed

+218
-5
lines changed

7 files changed

+218
-5
lines changed

‎dbt-redshift/.changes/unreleased/Features-20250710-110711.yaml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
kind: Features
2+
body: Support fine grained controls for tcp keepalive settings
3+
time: 2025-07-10T11:07:11.34467-07:00
4+
custom:
5+
Author: colin-rogers-dbt luis-garza-dev
6+
Issue: "876"

‎dbt-redshift/pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ dependencies = [
2828
"dbt-postgres>=1.8,<1.10",
2929
# dbt-redshift depends deeply on this package. it does not follow SemVer, therefore there have been breaking changes in previous patch releases
3030
# Pin to the patch or minor version, and bump in each new minor version of dbt-redshift.
31-
"redshift-connector>=2.1.3,<2.2",
31+
"redshift-connector>=2.1.8,<2.2",
3232
# add dbt-core to ensure backwards compatibility of installation, this is not a functional dependency
3333
"dbt-core>=1.8.0b3",
3434
# installed via dbt-core but referenced directly; don't pin to avoid version conflicts with dbt-core

‎dbt-redshift/src/dbt/adapters/redshift/connections.py‎

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -180,6 +180,12 @@ class RedshiftCredentials(Credentials):
180180
metadata={"description": "If using IAM auth, the AWS account id"},
181181
)
182182

183+
# TCP keepalive settings
184+
tcp_keepalive: Optional[bool] = True
185+
tcp_keepalive_idle: Optional[int] = None
186+
tcp_keepalive_interval: Optional[int] = None
187+
tcp_keepalive_count: Optional[int] = None
188+
183189
_ALIASES = {"dbname": "database", "pass": "password"}
184190

185191
@property
@@ -212,6 +218,10 @@ def _connection_keys(self):
212218
"is_serverless",
213219
"serverless_work_group",
214220
"serverless_acct_id",
221+
"tcp_keepalive",
222+
"tcp_keepalive_idle",
223+
"tcp_keepalive_interval",
224+
"tcp_keepalive_count",
215225
)
216226

217227
@property
@@ -243,7 +253,9 @@ def __base_kwargs(credentials) -> Dict[str, Any]:
243253
redshift_ssl_config: Dict[str, Any] = RedshiftSSLConfig.parse(
244254
credentials.sslmode
245255
).to_dict()
246-
return {
256+
257+
# Base connection parameters
258+
base_kwargs = {
247259
"host": credentials.host,
248260
"port": int(credentials.port) if credentials.port else 5439,
249261
"database": credentials.database,
@@ -255,6 +267,18 @@ def __base_kwargs(credentials) -> Dict[str, Any]:
255267
**redshift_ssl_config,
256268
}
257269

270+
# Add TCP keepalive parameters if enabled
271+
if credentials.tcp_keepalive:
272+
base_kwargs["tcp_keepalive"] = True
273+
if credentials.tcp_keepalive_idle is not None:
274+
base_kwargs["tcp_keepalive_idle"] = credentials.tcp_keepalive_idle
275+
if credentials.tcp_keepalive_interval is not None:
276+
base_kwargs["tcp_keepalive_interval"] = credentials.tcp_keepalive_interval
277+
if credentials.tcp_keepalive_count is not None:
278+
base_kwargs["tcp_keepalive_count"] = credentials.tcp_keepalive_count
279+
280+
return base_kwargs
281+
258282
def __iam_kwargs(credentials) -> Dict[str, Any]:
259283

260284
# iam True except for identity center methods

‎dbt-redshift/tests/functional/conftest.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,15 @@ def dbt_profile_target():
99
return {
1010
"type": "redshift",
1111
"host": os.getenv("REDSHIFT_TEST_HOST"),
12-
"port": int(os.getenv("REDSHIFT_TEST_PORT")),
12+
"port": int(os.getenv("REDSHIFT_TEST_PORT", "5439")),
1313
"dbname": os.getenv("REDSHIFT_TEST_DBNAME"),
1414
"user": os.getenv("REDSHIFT_TEST_USER"),
1515
"pass": os.getenv("REDSHIFT_TEST_PASS"),
1616
"region": os.getenv("REDSHIFT_TEST_REGION"),
1717
"threads": 1,
1818
"retries": 6,
19+
"tcp_keepalive": True,
20+
"tcp_keepalive_idle": 200,
21+
"tcp_keepalive_interval": 200,
22+
"tcp_keepalive_count": 5,
1923
}

‎dbt-redshift/tests/unit/test_auth_method.py‎

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717

1818
DEFAULT_SSL_CONFIG = RedshiftSSLConfig().to_dict()
19+
DEFAULT_TCP_KEEPALIVE_CONFIG = {"tcp_keepalive": True}
1920

2021

2122
class AuthMethod(TestCase):
@@ -61,9 +62,9 @@ def adapter(self):
6162
class TestInvalidMethod(AuthMethod):
6263
def test_invalid_auth_method(self):
6364
# we have to set method this way, otherwise it won't validate
64-
self.config.credentials.method = "badmethod"
65+
self.config.credentials.method = "badmethod" # type: ignore
6566
with self.assertRaises(FailedToConnectError) as context:
66-
connect_method_factory = get_connection_method(self.config.credentials)
67+
connect_method_factory = get_connection_method(self.config.credentials) # type: ignore
6768
connect_method_factory.get_connect_method()
6869
self.assertTrue("badmethod" in context.exception.msg)
6970

@@ -144,6 +145,7 @@ def test_default(self):
144145
region=None,
145146
is_serverless=False,
146147
**DEFAULT_SSL_CONFIG,
148+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
147149
)
148150

149151
@mock.patch("redshift_connector.connect", MagicMock())
@@ -164,6 +166,7 @@ def test_explicit_auth_method(self):
164166
timeout=None,
165167
is_serverless=False,
166168
**DEFAULT_SSL_CONFIG,
169+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
167170
)
168171

169172
def test_database_verification_is_case_insensitive(self):
@@ -258,6 +261,7 @@ def test_default(self):
258261
serverless_work_group=None,
259262
serverless_acct_id=None,
260263
**DEFAULT_SSL_CONFIG,
264+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
261265
)
262266

263267
@mock.patch("redshift_connector.connect", MagicMock())
@@ -289,6 +293,7 @@ def test_profile(self):
289293
serverless_work_group=None,
290294
serverless_acct_id=None,
291295
**DEFAULT_SSL_CONFIG,
296+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
292297
)
293298

294299
@mock.patch("redshift_connector.connect", MagicMock())
@@ -321,6 +326,7 @@ def test_explicit(self):
321326
serverless_work_group=None,
322327
serverless_acct_id=None,
323328
**DEFAULT_SSL_CONFIG,
329+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
324330
)
325331

326332
@mock.patch("redshift_connector.connect", MagicMock())
@@ -357,6 +363,7 @@ def test_explicit_workgroup_name(self):
357363
serverless_work_group="my_workgroup",
358364
serverless_acct_id="0123456789",
359365
**DEFAULT_SSL_CONFIG,
366+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
360367
)
361368

362369

@@ -389,6 +396,7 @@ def test_profile_default_region(self):
389396
serverless_work_group=None,
390397
serverless_acct_id=None,
391398
**DEFAULT_SSL_CONFIG,
399+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
392400
)
393401

394402
@mock.patch("redshift_connector.connect", MagicMock())
@@ -420,6 +428,7 @@ def test_profile_explicit_serverless(self):
420428
serverless_work_group=None,
421429
serverless_acct_id=None,
422430
**DEFAULT_SSL_CONFIG,
431+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
423432
)
424433

425434
@mock.patch("redshift_connector.connect", MagicMock())
@@ -451,6 +460,7 @@ def test_profile_explicit_region(self):
451460
serverless_work_group=None,
452461
serverless_acct_id=None,
453462
**DEFAULT_SSL_CONFIG,
463+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
454464
)
455465

456466
@mock.patch("redshift_connector.connect", MagicMock())
@@ -478,6 +488,7 @@ def test_profile_invalid_serverless(self):
478488
port=5439,
479489
timeout=None,
480490
**DEFAULT_SSL_CONFIG,
491+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
481492
)
482493
self.assertTrue("'host' must be provided" in context.exception.msg)
483494

@@ -518,6 +529,7 @@ def test_default(self):
518529
serverless_work_group=None,
519530
serverless_acct_id=None,
520531
**DEFAULT_SSL_CONFIG,
532+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
521533
)
522534

523535
@mock.patch("redshift_connector.connect", MagicMock())
@@ -548,6 +560,7 @@ def test_profile(self):
548560
serverless_work_group=None,
549561
serverless_acct_id=None,
550562
**DEFAULT_SSL_CONFIG,
563+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
551564
)
552565

553566

@@ -582,6 +595,7 @@ def test_profile_default_region(self):
582595
serverless_work_group=None,
583596
serverless_acct_id=None,
584597
**DEFAULT_SSL_CONFIG,
598+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
585599
)
586600

587601
@mock.patch("redshift_connector.connect", MagicMock())
@@ -613,6 +627,7 @@ def test_profile_ignore_cluster(self):
613627
serverless_work_group=None,
614628
serverless_acct_id=None,
615629
**DEFAULT_SSL_CONFIG,
630+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
616631
)
617632

618633
@mock.patch("redshift_connector.connect", MagicMock())
@@ -645,6 +660,7 @@ def test_profile_explicit_region(self):
645660
serverless_work_group=None,
646661
serverless_acct_id=None,
647662
**DEFAULT_SSL_CONFIG,
663+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
648664
)
649665

650666
@mock.patch("redshift_connector.connect", MagicMock())
@@ -676,6 +692,7 @@ def test_profile_invalid_serverless(self):
676692
serverless_work_group=None,
677693
serverless_acct_id=None,
678694
**DEFAULT_SSL_CONFIG,
695+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
679696
)
680697
self.assertTrue("'host' must be provided" in context.exception.msg)
681698

@@ -716,6 +733,7 @@ def test_profile_idc_browser_all_fields(self):
716733
idc_region="us-east-1",
717734
issuer_url="https://identitycenter.amazonaws.com/ssoins-randomchars",
718735
listen_port=1111,
736+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
719737
)
720738

721739
@mock.patch("redshift_connector.connect", MagicMock())
@@ -750,6 +768,7 @@ def test_profile_idc_browser_required_fields_only(self):
750768
idc_client_display_name="Amazon Redshift driver",
751769
idc_region="us-east-1",
752770
issuer_url="https://identitycenter.amazonaws.com/ssoins-randomchars",
771+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
753772
)
754773

755774
def test_invalid_adapter_missing_fields(self):
@@ -779,6 +798,7 @@ def test_invalid_adapter_missing_fields(self):
779798
listen_port=1111,
780799
idp_response_timeout=60,
781800
idc_client_display_name="my display",
801+
**DEFAULT_TCP_KEEPALIVE_CONFIG,
782802
)
783803

784804
assert (

0 commit comments

Comments
 (0)