Skip to content

Commit f096b89

Browse files
ericdeloreficentkathole
authored andcommitted
fix: Stop the Couchbase online store reporting failed writes as written
online_write_batch logged and skipped any upsert that raised, and still called progress for it, so materialization succeeded with rows missing from the online store. update() did the same when a scope or collection could not be created. The generic handlers now log and re-raise; the already-exists handlers are unchanged. Signed-off-by: Eric Del Orefice <114540347+ericdelorefice@users.noreply.github.com>
1 parent 8d2773c commit f096b89

2 files changed

Lines changed: 122 additions & 0 deletions

File tree

‎sdk/python/feast/infra/online_stores/couchbase_online_store/couchbase.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,7 @@ def online_write_batch(
143143
) # Upsert the document
144144
except Exception as e:
145145
logger.exception(f"Error upserting document {document_id}: {e}")
146+
raise
146147

147148
if progress:
148149
progress(1)
@@ -255,6 +256,7 @@ def update(
255256
logger.error(f"Scope {scope_name} already exists")
256257
except Exception as e:
257258
logger.error(f"Error creating scope {scope_name}: {e}")
259+
raise
258260

259261
# Check and create collection
260262
try:
@@ -268,6 +270,7 @@ def update(
268270
)
269271
except Exception as e:
270272
logger.error(f"Error creating collection {collection_name}: {e}")
273+
raise
271274

272275
def teardown(
273276
self,
Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
from datetime import datetime, timezone
2+
from unittest.mock import MagicMock, patch
3+
4+
import pytest
5+
6+
pytest.importorskip("couchbase", reason="couchbase not installed")
7+
8+
from couchbase.exceptions import ( # noqa: E402
9+
AuthenticationException,
10+
CollectionAlreadyExistsException,
11+
ScopeAlreadyExistsException,
12+
TimeoutException,
13+
)
14+
15+
from feast.infra.online_stores.couchbase_online_store.couchbase import ( # noqa: E402
16+
CouchbaseOnlineStore,
17+
CouchbaseOnlineStoreConfig,
18+
)
19+
from feast.protos.feast.types.EntityKey_pb2 import (
20+
EntityKey as EntityKeyProto, # noqa: E402
21+
)
22+
from feast.protos.feast.types.Value_pb2 import Value as ValueProto # noqa: E402
23+
from feast.repo_config import RepoConfig # noqa: E402
24+
25+
pytestmark = pytest.mark.filterwarnings("ignore::RuntimeWarning")
26+
27+
28+
@pytest.fixture
29+
def config():
30+
return RepoConfig(
31+
registry="dummy_registry.db",
32+
project="test_project",
33+
provider="local",
34+
entity_key_serialization_version=3,
35+
online_store=CouchbaseOnlineStoreConfig(),
36+
)
37+
38+
39+
@pytest.fixture
40+
def table():
41+
table = MagicMock()
42+
table.name = "driver_stats"
43+
return table
44+
45+
46+
def _rows(n):
47+
now = datetime.now(tz=timezone.utc)
48+
return [
49+
(
50+
EntityKeyProto(
51+
join_keys=["driver_id"], entity_values=[ValueProto(int64_val=i)]
52+
),
53+
{"trips": ValueProto(int64_val=i * 10)},
54+
now,
55+
None,
56+
)
57+
for i in range(n)
58+
]
59+
60+
61+
def _store_with(collection):
62+
store = CouchbaseOnlineStore()
63+
store.bucket = MagicMock()
64+
return store, patch.object(store, "_get_conn", return_value=collection)
65+
66+
67+
def test_write_batch_raises_when_upsert_fails(config, table):
68+
collection = MagicMock()
69+
collection.upsert.side_effect = TimeoutException("upsert timed out")
70+
progress = MagicMock()
71+
store, conn = _store_with(collection)
72+
73+
with conn, pytest.raises(TimeoutException):
74+
store.online_write_batch(config, table, _rows(3), progress)
75+
76+
# The failed row is not reported as written.
77+
progress.assert_not_called()
78+
79+
80+
def test_write_batch_writes_and_reports_every_row(config, table):
81+
collection = MagicMock()
82+
progress = MagicMock()
83+
store, conn = _store_with(collection)
84+
85+
with conn:
86+
store.online_write_batch(config, table, _rows(3), progress)
87+
88+
assert collection.upsert.call_count == 3
89+
assert progress.call_count == 3
90+
91+
92+
def test_update_raises_when_scope_cannot_be_created(config, table):
93+
store, conn = _store_with(MagicMock())
94+
store.bucket.collections.return_value.create_scope.side_effect = (
95+
AuthenticationException("not permitted")
96+
)
97+
98+
with conn, pytest.raises(AuthenticationException):
99+
store.update(config, [], [table], [], [], partial=False)
100+
101+
102+
def test_update_raises_when_collection_cannot_be_created(config, table):
103+
store, conn = _store_with(MagicMock())
104+
store.bucket.collections.return_value.create_collection.side_effect = (
105+
TimeoutException("create timed out")
106+
)
107+
108+
with conn, pytest.raises(TimeoutException):
109+
store.update(config, [], [table], [], [], partial=False)
110+
111+
112+
def test_update_accepts_scope_and_collection_that_already_exist(config, table):
113+
store, conn = _store_with(MagicMock())
114+
manager = store.bucket.collections.return_value
115+
manager.create_scope.side_effect = ScopeAlreadyExistsException("exists")
116+
manager.create_collection.side_effect = CollectionAlreadyExistsException("exists")
117+
118+
with conn:
119+
store.update(config, [], [table], [], [], partial=False)

0 commit comments

Comments
 (0)