Skip to content

Commit dff8dc6

Browse files
Pigbibicodex
andcommitted
fix: make Firstrade live claim atomic in GCS
Co-Authored-By: Codex <noreply@openai.com>
1 parent f07c1f0 commit dff8dc6

2 files changed

Lines changed: 67 additions & 1 deletion

File tree

application/state_persistence.py

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@ def _object_store():
2929
class GcsStateStore:
3030
bucket: str
3131
prefix: str = "firstrade-platform"
32+
client_factory: Callable[..., Any] | None = None
3233

3334
def __post_init__(self) -> None:
3435
if not str(self.bucket or "").strip():
@@ -74,8 +75,27 @@ def create_json(self, key: str, payload: dict[str, Any]) -> bool:
7475
"""Atomically create a JSON object; return False if it already exists."""
7576
uri = self._object_uri(key)
7677
data = json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
78+
bucket_name, separator, object_name = uri[5:].partition("/")
79+
if not separator or not bucket_name or not object_name:
80+
raise StatePersistenceError(f"Invalid GCS state URI for {key}")
7781
try:
78-
return bool(_object_store().create_text(uri, data, content_type="application/json"))
82+
from google.api_core.exceptions import Conflict, PreconditionFailed
83+
from google.cloud import storage
84+
85+
factory = self.client_factory or storage.Client
86+
try:
87+
client = factory()
88+
except TypeError:
89+
client = factory(project=None)
90+
blob = client.bucket(bucket_name).blob(object_name)
91+
blob.upload_from_string(
92+
data,
93+
content_type="application/json",
94+
if_generation_match=0,
95+
)
96+
return True
97+
except (Conflict, PreconditionFailed):
98+
return False
7999
except Exception as exc:
80100
raise StatePersistenceError(f"GCS atomic create failed for {key}: {exc}") from exc
81101

tests/test_strategy_run_claim.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
claim_live_strategy_run,
33
strategy_run_claim_key,
44
)
5+
from application.state_persistence import GcsStateStore
56

67

78
class AtomicFakeStore:
@@ -31,3 +32,48 @@ def test_live_claim_is_create_only_and_permanent():
3132
)
3233
assert store.payloads[key]["stage"] == "PENDING_SUBMISSION"
3334
assert store.payloads[key]["no_order_submitted"] is True
35+
36+
37+
def test_gcs_state_store_create_json_uses_generation_zero_precondition():
38+
class FakeBlob:
39+
def __init__(self):
40+
self.payload = None
41+
self.content_type = None
42+
self.if_generation_match = None
43+
44+
def upload_from_string(self, payload, *, content_type, if_generation_match):
45+
self.payload = payload
46+
self.content_type = content_type
47+
self.if_generation_match = if_generation_match
48+
49+
class FakeBucket:
50+
def __init__(self, blob):
51+
self._blob = blob
52+
self.name = ""
53+
self.object_name = ""
54+
55+
def blob(self, object_name):
56+
self.object_name = object_name
57+
return self._blob
58+
59+
class FakeClient:
60+
def __init__(self, bucket):
61+
self._bucket = bucket
62+
63+
def bucket(self, name):
64+
self._bucket.name = name
65+
return self._bucket
66+
67+
blob = FakeBlob()
68+
bucket = FakeBucket(blob)
69+
store = GcsStateStore(
70+
bucket="state-bucket",
71+
prefix="runtime",
72+
client_factory=lambda: FakeClient(bucket),
73+
)
74+
75+
assert store.create_json("claims/run.json", {"state": "claimed"}) is True
76+
assert bucket.name == "state-bucket"
77+
assert bucket.object_name == "runtime/claims/run.json"
78+
assert blob.content_type == "application/json"
79+
assert blob.if_generation_match == 0

0 commit comments

Comments
 (0)