Skip to content

Commit 51ce982

Browse files
casaar97ntkathole
casaar97
authored and
ntkathole
committed
feat: add plan() support to DynamoDBOnlineStore
DynamoDBOnlineStore has no plan() override, so it inherits OnlineStore.plan()'s no-op default: `feast plan` never reports any DynamoDB infrastructure changes, unlike SqliteOnlineStore and MilvusOnlineStore which both implement it. Adds DynamoDBTable (InfraObject) and DynamoDBOnlineStore.plan(), using the InfraObject proto's CustomInfra field (protos/feast/core/InfraObject.proto), which exists specifically so online stores can add InfraObject support without changes to the core proto -- no other in-tree store uses it yet. Note this only affects `feast plan`'s reporting: `feast apply`'s diff-based path (FeatureStore._should_use_plan(), which would call InfraObject.update()/teardown()) is gated to the local/sqlite provider only, so DynamoDBOnlineStore.update()/teardown() -- which already perform the real table creation/deletion -- are unaffected. Uses the corrected feature-view-list pattern (see #6658 / #6659): FeatureView.from_proto() and StreamFeatureView.from_proto() applied to their respective proto lists, not one applied to both. Signed-off-by: Carlos Sánchez <carlos.sancheza@cabify.com>
1 parent 104ad10 commit 51ce982

2 files changed

Lines changed: 245 additions & 0 deletions

File tree

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

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import asyncio
1515
import contextlib
1616
import itertools
17+
import json
1718
import logging
1819
from collections import OrderedDict, defaultdict
1920
from concurrent.futures import ThreadPoolExecutor
@@ -24,13 +25,17 @@
2425
from pydantic import StrictBool, StrictStr
2526

2627
from feast import Entity, FeatureView, utils
28+
from feast.infra.infra_object import InfraObject
2729
from feast.infra.online_stores.helpers import compute_entity_id, compute_versioned_name
2830
from feast.infra.online_stores.online_store import OnlineStore
2931
from feast.infra.supported_async_methods import SupportedAsyncMethods
3032
from feast.infra.utils.aws_utils import dynamo_write_items_async
33+
from feast.protos.feast.core.InfraObject_pb2 import InfraObject as InfraObjectProto
34+
from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto
3135
from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto
3236
from feast.protos.feast.types.Value_pb2 import Value as ValueProto
3337
from feast.repo_config import FeastConfigBaseModel, RepoConfig
38+
from feast.stream_feature_view import StreamFeatureView
3439
from feast.utils import get_user_agent
3540

3641
try:
@@ -379,6 +384,31 @@ def update(
379384
_get_table_name(online_config, config, table_to_delete),
380385
)
381386

387+
def plan(
388+
self, config: RepoConfig, desired_registry_proto: RegistryProto
389+
) -> List[InfraObject]:
390+
online_config = config.online_store
391+
assert isinstance(online_config, DynamoDBOnlineStoreConfig)
392+
393+
# feature_views and stream_feature_views are distinct proto types
394+
# (FeatureViewProto vs StreamFeatureViewProto); each needs its
395+
# matching from_proto(), not one applied to both indiscriminately.
396+
views = [
397+
FeatureView.from_proto(view)
398+
for view in desired_registry_proto.feature_views
399+
] + [
400+
StreamFeatureView.from_proto(view)
401+
for view in desired_registry_proto.stream_feature_views
402+
]
403+
return [
404+
DynamoDBTable(
405+
name=_get_table_name(online_config, config, view),
406+
region=online_config.region,
407+
endpoint_url=online_config.endpoint_url,
408+
)
409+
for view in views
410+
]
411+
382412
def teardown(
383413
self,
384414
config: RepoConfig,
@@ -1222,6 +1252,68 @@ def _get_table_name(
12221252
)
12231253

12241254

1255+
DYNAMODB_INFRA_OBJECT_CLASS_TYPE = "feast.infra.online_stores.dynamodb.DynamoDBTable"
1256+
1257+
1258+
class DynamoDBTable(InfraObject):
1259+
"""
1260+
A DynamoDB table managed by Feast, reported by DynamoDBOnlineStore.plan()
1261+
so `feast plan` can show DynamoDB table changes.
1262+
1263+
Uses the InfraObject proto's CustomInfra field rather than a dedicated
1264+
proto message, since that field exists precisely to let online stores
1265+
add InfraObject support without changing the core InfraObject proto.
1266+
1267+
Note: `feast apply` does not call update()/teardown() on this object --
1268+
DynamoDBOnlineStore.update()/teardown() perform the actual table
1269+
creation and deletion directly (see FeatureStore._should_use_plan(),
1270+
which gates the new diff-based apply path to the local/sqlite provider
1271+
only). update()/teardown() here are no-ops that satisfy InfraObject's
1272+
abstract interface.
1273+
"""
1274+
1275+
def __init__(self, name: str, region: str, endpoint_url: Optional[str] = None):
1276+
super().__init__(name)
1277+
self.region = region
1278+
self.endpoint_url = endpoint_url
1279+
1280+
def to_infra_object_proto(self) -> InfraObjectProto:
1281+
return InfraObjectProto(
1282+
infra_object_class_type=DYNAMODB_INFRA_OBJECT_CLASS_TYPE,
1283+
custom_infra=InfraObjectProto.CustomInfra(field=self.to_proto()),
1284+
)
1285+
1286+
def to_proto(self) -> bytes:
1287+
return json.dumps(
1288+
{
1289+
"name": self.name,
1290+
"region": self.region,
1291+
"endpoint_url": self.endpoint_url,
1292+
}
1293+
).encode("utf-8")
1294+
1295+
@staticmethod
1296+
def from_infra_object_proto(
1297+
infra_object_proto: InfraObjectProto,
1298+
) -> "DynamoDBTable":
1299+
return DynamoDBTable.from_proto(infra_object_proto.custom_infra.field)
1300+
1301+
@staticmethod
1302+
def from_proto(serialized: bytes) -> "DynamoDBTable":
1303+
payload = json.loads(serialized.decode("utf-8"))
1304+
return DynamoDBTable(
1305+
name=payload["name"],
1306+
region=payload["region"],
1307+
endpoint_url=payload.get("endpoint_url"),
1308+
)
1309+
1310+
def update(self) -> None:
1311+
pass
1312+
1313+
def teardown(self) -> None:
1314+
pass
1315+
1316+
12251317
def _delete_table_idempotent(
12261318
dynamodb_resource,
12271319
table_name: str,

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

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
from datetime import timedelta
2+
3+
from feast.data_format import AvroFormat
4+
from feast.data_source import KafkaSource
5+
from feast.feature_view import FeatureView
6+
from feast.field import Field
7+
from feast.infra.infra_object import Infra
8+
from feast.infra.offline_stores.dask import DaskOfflineStoreConfig
9+
from feast.infra.offline_stores.file_source import FileSource
10+
from feast.infra.online_stores.dynamodb import (
11+
DynamoDBOnlineStore,
12+
DynamoDBOnlineStoreConfig,
13+
DynamoDBTable,
14+
)
15+
from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto
16+
from feast.repo_config import RepoConfig
17+
from feast.stream_feature_view import StreamFeatureView
18+
from feast.types import String
19+
20+
REGISTRY = "s3://test_registry/registry.db"
21+
PROJECT = "test_aws"
22+
PROVIDER = "aws"
23+
REGION = "us-west-2"
24+
25+
26+
def _repo_config() -> RepoConfig:
27+
return RepoConfig(
28+
registry=REGISTRY,
29+
project=PROJECT,
30+
provider=PROVIDER,
31+
online_store=DynamoDBOnlineStoreConfig(region=REGION),
32+
offline_store=DaskOfflineStoreConfig(),
33+
entity_key_serialization_version=3,
34+
)
35+
36+
37+
def _feature_view(name: str) -> FeatureView:
38+
return FeatureView(
39+
name=name,
40+
entities=[],
41+
schema=[Field(name="value", dtype=String)],
42+
ttl=timedelta(days=1),
43+
online=True,
44+
source=FileSource(path="dummy.parquet", timestamp_field="event_timestamp"),
45+
)
46+
47+
48+
def _stream_feature_view(name: str) -> StreamFeatureView:
49+
return StreamFeatureView(
50+
name=name,
51+
entities=[],
52+
schema=[Field(name="value", dtype=String)],
53+
source=KafkaSource(
54+
name="dummy_kafka",
55+
timestamp_field="event_timestamp",
56+
message_format=AvroFormat(""),
57+
kafka_bootstrap_servers="localhost:9092",
58+
topic="dummy_topic",
59+
batch_source=FileSource(
60+
path="dummy.parquet", timestamp_field="event_timestamp"
61+
),
62+
),
63+
)
64+
65+
66+
class TestDynamoDBOnlineStorePlan:
67+
"""DynamoDBOnlineStore previously had no plan() override at all, so
68+
`feast plan` reported no DynamoDB infrastructure changes. This adds one,
69+
following the InfraObject.CustomInfra extension point (see
70+
protos/feast/core/InfraObject.proto) used by no other in-tree store yet."""
71+
72+
def test_plan_returns_one_table_per_feature_view(self):
73+
config = _repo_config()
74+
registry_proto = RegistryProto()
75+
registry_proto.feature_views.append(_feature_view("view_a").to_proto())
76+
registry_proto.feature_views.append(_feature_view("view_b").to_proto())
77+
78+
infra_objects = DynamoDBOnlineStore().plan(config, registry_proto)
79+
80+
assert sorted(o.name for o in infra_objects) == [
81+
f"{PROJECT}.view_a",
82+
f"{PROJECT}.view_b",
83+
]
84+
assert all(isinstance(o, DynamoDBTable) for o in infra_objects)
85+
assert all(o.region == REGION for o in infra_objects)
86+
87+
def test_plan_includes_stream_feature_views(self):
88+
"""Regression guard: FeatureView.from_proto() is @typechecked and
89+
only accepts a FeatureViewProto, so it can't be applied to
90+
stream_feature_views (StreamFeatureViewProto) too -- each list needs
91+
its matching class (see the related fix in sqlite.py's plan())."""
92+
config = _repo_config()
93+
registry_proto = RegistryProto()
94+
registry_proto.stream_feature_views.append(
95+
_stream_feature_view("driver_dropoffs_stream").to_proto()
96+
)
97+
98+
infra_objects = DynamoDBOnlineStore().plan(config, registry_proto)
99+
100+
assert [o.name for o in infra_objects] == [f"{PROJECT}.driver_dropoffs_stream"]
101+
102+
def test_plan_includes_batch_and_stream_feature_views_together(self):
103+
config = _repo_config()
104+
registry_proto = RegistryProto()
105+
registry_proto.feature_views.append(_feature_view("batch_view").to_proto())
106+
registry_proto.stream_feature_views.append(
107+
_stream_feature_view("driver_dropoffs_stream").to_proto()
108+
)
109+
110+
infra_objects = DynamoDBOnlineStore().plan(config, registry_proto)
111+
112+
assert sorted(o.name for o in infra_objects) == [
113+
f"{PROJECT}.batch_view",
114+
f"{PROJECT}.driver_dropoffs_stream",
115+
]
116+
117+
def test_plan_empty_registry_produces_empty_plan(self):
118+
config = _repo_config()
119+
registry_proto = RegistryProto()
120+
121+
assert DynamoDBOnlineStore().plan(config, registry_proto) == []
122+
123+
124+
class TestDynamoDBTableProtoRoundTrip:
125+
"""This is the mechanism `feast plan` relies on: Infra gets serialized to
126+
proto and deserialized back via the dotted class_type path, not a
127+
hardcoded per-store branch."""
128+
129+
def test_infra_object_survives_proto_round_trip(self):
130+
original = DynamoDBTable(
131+
name=f"{PROJECT}.view_a", region=REGION, endpoint_url=None
132+
)
133+
infra = Infra(infra_objects=[original])
134+
135+
restored = Infra.from_proto(infra.to_proto())
136+
137+
assert len(restored.infra_objects) == 1
138+
restored_table = restored.infra_objects[0]
139+
assert isinstance(restored_table, DynamoDBTable)
140+
assert restored_table.name == f"{PROJECT}.view_a"
141+
assert restored_table.region == REGION
142+
143+
def test_round_trip_preserves_endpoint_url(self):
144+
original = DynamoDBTable(
145+
name=f"{PROJECT}.view_a",
146+
region=REGION,
147+
endpoint_url="http://localhost:8000",
148+
)
149+
infra = Infra(infra_objects=[original])
150+
151+
restored = Infra.from_proto(infra.to_proto())
152+
153+
assert restored.infra_objects[0].endpoint_url == "http://localhost:8000"

0 commit comments

Comments
 (0)