Skip to content

Commit 77c9f13

Browse files
committed
feat: DynamoDBTable support endpoint_url
Signed-off-by: Miguel Trejo <[email protected]>
1 parent 4b9fcc6 commit 77c9f13

2 files changed

Lines changed: 75 additions & 10 deletions

File tree

‎sdk/python/feast/infra/online_stores/dynamodb.py‎

Lines changed: 23 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -342,13 +342,20 @@ class DynamoDBTable(InfraObject):
342342
Attributes:
343343
name: The name of the table.
344344
region: The region of the table.
345+
endpoint_url: Local DynamoDB Endpoint Url.
346+
_dynamodb_client: Boto3 DynamoDB client.
347+
_dynamodb_resource: Boto3 DynamoDB resource.
345348
"""
346349

347350
region: str
351+
endpoint_url = None
352+
_dynamodb_client = None
353+
_dynamodb_resource = None
348354

349-
def __init__(self, name: str, region: str):
355+
def __init__(self, name: str, region: str, endpoint_url: Optional[str] = None):
350356
super().__init__(name)
351357
self.region = region
358+
self.endpoint_url = endpoint_url
352359

353360
def to_infra_object_proto(self) -> InfraObjectProto:
354361
dynamodb_table_proto = self.to_proto()
@@ -377,12 +384,8 @@ def from_proto(dynamodb_table_proto: DynamoDBTableProto) -> Any:
377384
)
378385

379386
def update(self):
380-
dynamodb_client = _initialize_dynamodb_client(
381-
region=self.region, endpoint_url=None
382-
)
383-
dynamodb_resource = _initialize_dynamodb_resource(
384-
region=self.region, endpoint_url=None
385-
)
387+
dynamodb_client = self._get_dynamodb_client(self.region, self.endpoint_url)
388+
dynamodb_resource = self._get_dynamodb_resource(self.region, self.endpoint_url)
386389

387390
try:
388391
dynamodb_resource.create_table(
@@ -403,7 +406,17 @@ def update(self):
403406
dynamodb_client.get_waiter("table_exists").wait(TableName=f"{self.name}")
404407

405408
def teardown(self):
406-
dynamodb_resource = _initialize_dynamodb_resource(
407-
region=self.region, endpoint_url=None
408-
)
409+
dynamodb_resource = self._get_dynamodb_resource(self.region, self.endpoint_url)
409410
_delete_table_idempotent(dynamodb_resource, self.name)
411+
412+
def _get_dynamodb_client(self, region: str, endpoint_url: Optional[str] = None):
413+
if self._dynamodb_client is None:
414+
self._dynamodb_client = _initialize_dynamodb_client(region, endpoint_url)
415+
return self._dynamodb_client
416+
417+
def _get_dynamodb_resource(self, region: str, endpoint_url: Optional[str] = None):
418+
if self._dynamodb_resource is None:
419+
self._dynamodb_resource = _initialize_dynamodb_resource(
420+
region, endpoint_url
421+
)
422+
return self._dynamodb_resource

‎sdk/python/tests/unit/infra/online_store/test_dynamodb_online_store.py‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
from feast.infra.online_stores.dynamodb import (
88
DynamoDBOnlineStore,
99
DynamoDBOnlineStoreConfig,
10+
DynamoDBTable,
1011
)
1112
from feast.repo_config import RepoConfig
1213
from tests.utils.online_store_utils import (
@@ -50,6 +51,18 @@ def test_online_store_config_default():
5051
assert dynamodb_store_config.table_name_template == "{project}.{table_name}"
5152

5253

54+
def test_dynamodb_table_default_params():
55+
"""Test DynamoDBTable default parameters."""
56+
tbl_name = "dynamodb-test"
57+
aws_region = "us-west-2"
58+
dynamodb_table = DynamoDBTable(tbl_name, aws_region)
59+
assert dynamodb_table.name == tbl_name
60+
assert dynamodb_table.region == aws_region
61+
assert dynamodb_table.endpoint_url is None
62+
assert dynamodb_table._dynamodb_client is None
63+
assert dynamodb_table._dynamodb_resource is None
64+
65+
5366
def test_online_store_config_custom_params():
5467
"""Test DynamoDBOnlineStoreConfig custom parameters."""
5568
aws_region = "us-west-2"
@@ -72,6 +85,19 @@ def test_online_store_config_custom_params():
7285
assert dynamodb_store_config.table_name_template == table_name_template
7386

7487

88+
def test_dynamodb_table_custom_params():
89+
"""Test DynamoDBTable custom parameters."""
90+
tbl_name = "dynamodb-test"
91+
aws_region = "us-west-2"
92+
endpoint_url = "http://localhost:8000"
93+
dynamodb_table = DynamoDBTable(tbl_name, aws_region, endpoint_url)
94+
assert dynamodb_table.name == tbl_name
95+
assert dynamodb_table.region == aws_region
96+
assert dynamodb_table.endpoint_url == endpoint_url
97+
assert dynamodb_table._dynamodb_client is None
98+
assert dynamodb_table._dynamodb_resource is None
99+
100+
75101
def test_online_store_config_dynamodb_client():
76102
"""Test DynamoDBOnlineStoreConfig configure DynamoDB client with endpoint_url."""
77103
aws_region = "us-west-2"
@@ -87,6 +113,19 @@ def test_online_store_config_dynamodb_client():
87113
assert dynamodb_client.meta.endpoint_url == endpoint_url
88114

89115

116+
def test_dynamodb_table_dynamodb_client():
117+
"""Test DynamoDBTable configure DynamoDB client with endpoint_url."""
118+
tbl_name = "dynamodb-test"
119+
aws_region = "us-west-2"
120+
endpoint_url = "http://localhost:8000"
121+
dynamodb_table = DynamoDBTable(tbl_name, aws_region, endpoint_url)
122+
dynamodb_client = dynamodb_table._get_dynamodb_client(
123+
dynamodb_table.region, dynamodb_table.endpoint_url
124+
)
125+
assert dynamodb_client.meta.region_name == aws_region
126+
assert dynamodb_client.meta.endpoint_url == endpoint_url
127+
128+
90129
def test_online_store_config_dynamodb_resource():
91130
"""Test DynamoDBOnlineStoreConfig configure DynamoDB Resource with endpoint_url."""
92131
aws_region = "us-west-2"
@@ -102,6 +141,19 @@ def test_online_store_config_dynamodb_resource():
102141
assert dynamodb_resource.meta.client.meta.endpoint_url == endpoint_url
103142

104143

144+
def test_dynamodb_table_dynamodb_resource():
145+
"""Test DynamoDBTable configure DynamoDB resource with endpoint_url."""
146+
tbl_name = "dynamodb-test"
147+
aws_region = "us-west-2"
148+
endpoint_url = "http://localhost:8000"
149+
dynamodb_table = DynamoDBTable(tbl_name, aws_region, endpoint_url)
150+
dynamodb_resource = dynamodb_table._get_dynamodb_resource(
151+
dynamodb_table.region, dynamodb_table.endpoint_url
152+
)
153+
assert dynamodb_resource.meta.client.meta.region_name == aws_region
154+
assert dynamodb_resource.meta.client.meta.endpoint_url == endpoint_url
155+
156+
105157
@mock_dynamodb2
106158
@pytest.mark.parametrize("n_samples", [5, 50, 100])
107159
def test_online_read(repo_config, n_samples):

0 commit comments

Comments
 (0)