Skip to content

Commit 3a787ac

Browse files
robhowleyrobhowley
authored
feat: Add feature view tags to dynamo tags (feast-dev#5291)
* add feature view tags to dynamo tags Signed-off-by: Rob Howley <howley.robert@gmail.com> * add tag update helper method Signed-off-by: Rob Howley <howley.robert@gmail.com> * fix tag change checks Signed-off-by: Rob Howley <howley.robert@gmail.com> * simplify and just untag all to re add Signed-off-by: Rob Howley <howley.robert@gmail.com> * dont do tag updates for new tables Signed-off-by: Rob Howley <howley.robert@gmail.com> * more granular tags take priority Signed-off-by: Rob Howley <howley.robert@gmail.com> --------- Signed-off-by: Rob Howley <howley.robert@gmail.com>
1 parent e236c49 commit 3a787ac

2 files changed

Lines changed: 159 additions & 18 deletions

File tree

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

Lines changed: 52 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
import contextlib
1616
import itertools
1717
import logging
18-
from collections import OrderedDict
18+
from collections import OrderedDict, defaultdict
1919
from datetime import datetime
2020
from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple, Union
2121

@@ -138,6 +138,38 @@ async def close(self):
138138
def async_supported(self) -> SupportedAsyncMethods:
139139
return SupportedAsyncMethods(read=True, write=True)
140140

141+
@staticmethod
142+
def _table_tags(online_config, table_instance) -> list[dict[str, str]]:
143+
table_instance_tags = table_instance.tags or {}
144+
online_tags = online_config.tags or {}
145+
146+
common_tags = [
147+
{"Key": key, "Value": table_instance_tags.get(key) or value}
148+
for key, value in online_tags.items()
149+
]
150+
table_tags = [
151+
{"Key": key, "Value": value}
152+
for key, value in table_instance_tags.items()
153+
if key not in online_tags
154+
]
155+
156+
return common_tags + table_tags
157+
158+
@staticmethod
159+
def _update_tags(dynamodb_client, table_name: str, new_tags: list[dict[str, str]]):
160+
table_arn = dynamodb_client.describe_table(TableName=table_name)["Table"][
161+
"TableArn"
162+
]
163+
current_tags = dynamodb_client.list_tags_of_resource(ResourceArn=table_arn)[
164+
"Tags"
165+
]
166+
if current_tags:
167+
remove_keys = [tag["Key"] for tag in current_tags]
168+
dynamodb_client.untag_resource(ResourceArn=table_arn, TagKeys=remove_keys)
169+
170+
if new_tags:
171+
dynamodb_client.tag_resource(ResourceArn=table_arn, Tags=new_tags)
172+
141173
def update(
142174
self,
143175
config: RepoConfig,
@@ -167,40 +199,43 @@ def update(
167199
online_config.endpoint_url,
168200
online_config.session_based_auth,
169201
)
170-
# Add Tags attribute to creation request only if configured to prevent
171-
# TagResource permission issues, even with an empty Tags array.
172-
kwargs = (
173-
{
174-
"Tags": [
175-
{"Key": key, "Value": value}
176-
for key, value in online_config.tags.items()
177-
]
178-
}
179-
if online_config.tags
180-
else {}
181-
)
202+
203+
do_tag_updates = defaultdict(bool)
182204
for table_instance in tables_to_keep:
205+
# Add Tags attribute to creation request only if configured to prevent
206+
# TagResource permission issues, even with an empty Tags array.
207+
table_tags = self._table_tags(online_config, table_instance)
208+
kwargs = {"Tags": table_tags} if table_tags else {}
209+
210+
table_name = _get_table_name(online_config, config, table_instance)
183211
try:
184212
dynamodb_resource.create_table(
185-
TableName=_get_table_name(online_config, config, table_instance),
213+
TableName=table_name,
186214
KeySchema=[{"AttributeName": "entity_id", "KeyType": "HASH"}],
187215
AttributeDefinitions=[
188216
{"AttributeName": "entity_id", "AttributeType": "S"}
189217
],
190218
BillingMode="PAY_PER_REQUEST",
191219
**kwargs,
192220
)
221+
193222
except ClientError as ce:
223+
do_tag_updates[table_name] = True
224+
194225
# If the table creation fails with ResourceInUseException,
195226
# it means the table already exists or is being created.
196227
# Otherwise, re-raise the exception
197228
if ce.response["Error"]["Code"] != "ResourceInUseException":
198229
raise
199230

200231
for table_instance in tables_to_keep:
201-
dynamodb_client.get_waiter("table_exists").wait(
202-
TableName=_get_table_name(online_config, config, table_instance)
203-
)
232+
table_name = _get_table_name(online_config, config, table_instance)
233+
dynamodb_client.get_waiter("table_exists").wait(TableName=table_name)
234+
# once table is confirmed to exist, update the tags.
235+
# tags won't be updated in the create_table call if the table already exists
236+
if do_tag_updates[table_name]:
237+
tags = self._table_tags(online_config, table_instance)
238+
self._update_tags(dynamodb_client, table_name, tags)
204239

205240
for table_to_delete in tables_to_delete:
206241
_delete_table_idempotent(

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

Lines changed: 107 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from copy import deepcopy
22
from dataclasses import dataclass
33
from datetime import datetime
4+
from typing import Optional
45

56
import boto3
67
import pytest
@@ -32,6 +33,12 @@
3233
@dataclass
3334
class MockFeatureView:
3435
name: str
36+
tags: Optional[dict[str, str]] = None
37+
38+
39+
@dataclass
40+
class MockOnlineConfig:
41+
tags: Optional[dict[str, str]] = None
3542

3643

3744
@pytest.fixture
@@ -209,6 +216,13 @@ def test_dynamodb_online_store_online_write_batch(
209216
assert [item[1] for item in stored_items] == list(features)
210217

211218

219+
def _get_tags(dynamodb_client, table_name):
220+
table_arn = dynamodb_client.describe_table(TableName=table_name)["Table"][
221+
"TableArn"
222+
]
223+
return dynamodb_client.list_tags_of_resource(ResourceArn=table_arn).get("Tags")
224+
225+
212226
@mock_dynamodb
213227
def test_dynamodb_online_store_update(repo_config, dynamodb_online_store):
214228
"""Test DynamoDBOnlineStore update method."""
@@ -222,7 +236,7 @@ def test_dynamodb_online_store_update(repo_config, dynamodb_online_store):
222236
dynamodb_online_store.update(
223237
config=repo_config,
224238
tables_to_delete=[MockFeatureView(name=db_table_delete_name)],
225-
tables_to_keep=[MockFeatureView(name=db_table_keep_name)],
239+
tables_to_keep=[MockFeatureView(name=db_table_keep_name, tags={"some": "tag"})],
226240
entities_to_delete=None,
227241
entities_to_keep=None,
228242
partial=None,
@@ -237,6 +251,98 @@ def test_dynamodb_online_store_update(repo_config, dynamodb_online_store):
237251
assert len(existing_tables) == 1
238252
assert existing_tables[0] == f"test_aws.{db_table_keep_name}"
239253

254+
assert _get_tags(dynamodb_client, existing_tables[0]) == [
255+
{"Key": "some", "Value": "tag"}
256+
]
257+
258+
259+
@mock_dynamodb
260+
def test_dynamodb_online_store_update_tags(repo_config, dynamodb_online_store):
261+
"""Test DynamoDBOnlineStore update method."""
262+
# create dummy table to update with new tags and tag values
263+
table_name = f"{TABLE_NAME}_keep_update_tags"
264+
create_test_table(PROJECT, table_name, REGION)
265+
266+
# add tags on update
267+
dynamodb_online_store.update(
268+
config=repo_config,
269+
tables_to_delete=[],
270+
tables_to_keep=[
271+
MockFeatureView(
272+
name=table_name, tags={"key1": "val1", "key2": "val2", "key3": "val3"}
273+
)
274+
],
275+
entities_to_delete=[],
276+
entities_to_keep=[],
277+
partial=None,
278+
)
279+
280+
# update tags
281+
dynamodb_online_store.update(
282+
config=repo_config,
283+
tables_to_delete=[],
284+
tables_to_keep=[
285+
MockFeatureView(
286+
name=table_name,
287+
tags={"key1": "new-val1", "key2": "val2", "key4": "val4"},
288+
)
289+
],
290+
entities_to_delete=[],
291+
entities_to_keep=[],
292+
partial=None,
293+
)
294+
295+
# check only db_table_keep_name exists
296+
dynamodb_client = dynamodb_online_store._get_dynamodb_client(REGION)
297+
existing_tables = dynamodb_client.list_tables().get("TableNames", None)
298+
299+
expected_tags = [
300+
{"Key": "key1", "Value": "new-val1"},
301+
{"Key": "key2", "Value": "val2"},
302+
{"Key": "key4", "Value": "val4"},
303+
]
304+
assert _get_tags(dynamodb_client, existing_tables[0]) == expected_tags
305+
306+
# and then remove all tags
307+
dynamodb_online_store.update(
308+
config=repo_config,
309+
tables_to_delete=[],
310+
tables_to_keep=[MockFeatureView(name=table_name, tags=None)],
311+
entities_to_delete=[],
312+
entities_to_keep=[],
313+
partial=None,
314+
)
315+
316+
assert _get_tags(dynamodb_client, existing_tables[0]) == []
317+
318+
319+
@mock_dynamodb
320+
@pytest.mark.parametrize(
321+
"global_tags, table_tags, expected",
322+
[
323+
(None, {"key": "val"}, [{"Key": "key", "Value": "val"}]),
324+
({"key": "val"}, None, [{"Key": "key", "Value": "val"}]),
325+
(
326+
{"key1": "val1"},
327+
{"key2": "val2"},
328+
[{"Key": "key1", "Value": "val1"}, {"Key": "key2", "Value": "val2"}],
329+
),
330+
(
331+
{"key": "val", "key2": "val2"},
332+
{"key": "new-val"},
333+
[{"Key": "key", "Value": "new-val"}, {"Key": "key2", "Value": "val2"}],
334+
),
335+
],
336+
)
337+
def test_dynamodb_online_store_tag_priority(
338+
global_tags, table_tags, expected, dynamodb_online_store
339+
):
340+
actual = dynamodb_online_store._table_tags(
341+
MockOnlineConfig(tags=global_tags),
342+
MockFeatureView(name="table", tags=table_tags),
343+
)
344+
assert actual == expected
345+
240346

241347
@mock_dynamodb
242348
def test_dynamodb_online_store_teardown(repo_config, dynamodb_online_store):

0 commit comments

Comments
 (0)