Skip to content

Commit a0273f4

Browse files
Yashwanth-Ranjan-SingaravelYashwanth-Ranjan-Singaravel
authored
A simple kafka lambda that interacts with a DB (#440)
* Initial experiment set up * Added instructions + simplifications to f.py
1 parent 1a54ec1 commit a0273f4

6 files changed

Lines changed: 394 additions & 0 deletions

File tree

‎examples/kafka-db-sum/f.py‎

Lines changed: 198 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,198 @@
1+
from flask import Flask, request, make_response, jsonify
2+
import os
3+
import random
4+
import psycopg2
5+
6+
app = Flask(__name__)
7+
8+
DATABASE_URL = os.environ.get("DATABASE_URL")
9+
10+
# Probability (0.0–1.0) that a transaction will fail between UPDATE and COMMIT.
11+
# Set to 0 for normal operation; raise to stress-test seek-back recovery.
12+
FAIL_PROBABILITY = float(os.environ.get("FAIL_PROBABILITY", "0"))
13+
_db_initialized = False
14+
15+
16+
def get_db():
17+
return psycopg2.connect(DATABASE_URL)
18+
19+
20+
def ensure_db():
21+
"""Create the running_sum table if it doesn't exist (runs once per sandbox)."""
22+
global _db_initialized
23+
if _db_initialized:
24+
return
25+
conn = get_db()
26+
try:
27+
with conn.cursor() as cur:
28+
cur.execute(
29+
"""
30+
CREATE TABLE IF NOT EXISTS running_sum (
31+
id INTEGER PRIMARY KEY DEFAULT 1,
32+
total BIGINT NOT NULL DEFAULT 0,
33+
last_offset BIGINT NOT NULL DEFAULT -1,
34+
message_count BIGINT NOT NULL DEFAULT 0,
35+
CHECK (id = 1)
36+
)
37+
"""
38+
)
39+
cur.execute(
40+
"""
41+
INSERT INTO running_sum (id, total, last_offset, message_count)
42+
VALUES (1, 0, -1, 0)
43+
ON CONFLICT (id) DO NOTHING
44+
"""
45+
)
46+
conn.commit()
47+
_db_initialized = True
48+
finally:
49+
conn.close()
50+
51+
52+
@app.route("/reset", methods=["POST"])
53+
def reset():
54+
"""Reset running_sum to zero so the demo can be re-run cleanly."""
55+
ensure_db()
56+
conn = get_db()
57+
try:
58+
with conn.cursor() as cur:
59+
cur.execute(
60+
"UPDATE running_sum SET total = 0, last_offset = -1, "
61+
"message_count = 0 WHERE id = 1"
62+
)
63+
conn.commit()
64+
return jsonify({"status": "reset"})
65+
finally:
66+
conn.close()
67+
68+
69+
@app.route("/", methods=["GET", "POST"])
70+
def handle():
71+
ensure_db()
72+
73+
# GET — return current state (useful for checking progress via HTTP)
74+
if request.method == "GET":
75+
conn = get_db()
76+
try:
77+
with conn.cursor() as cur:
78+
cur.execute(
79+
"SELECT total, last_offset, message_count "
80+
"FROM running_sum WHERE id = 1"
81+
)
82+
row = cur.fetchone()
83+
if row:
84+
return jsonify(
85+
{
86+
"running_sum": row[0],
87+
"last_offset": row[1],
88+
"message_count": row[2],
89+
}
90+
)
91+
return jsonify(
92+
{"running_sum": 0, "last_offset": -1, "message_count": 0}
93+
)
94+
finally:
95+
conn.close()
96+
97+
# POST — process a Kafka message containing a number
98+
offset = int(request.headers.get("X-Kafka-Offset", "-1"))
99+
topic = request.headers.get("X-Kafka-Topic", "unknown")
100+
partition = request.headers.get("X-Kafka-Partition", "unknown")
101+
102+
body = request.get_json(silent=True)
103+
104+
# Accept {"number": N}
105+
if isinstance(body, dict):
106+
number = body.get("number", 0)
107+
108+
conn = None
109+
try:
110+
conn = get_db()
111+
with conn.cursor() as cur:
112+
# Lock the row and read last processed offset
113+
cur.execute(
114+
"SELECT last_offset FROM running_sum WHERE id = 1 FOR UPDATE"
115+
)
116+
row = cur.fetchone()
117+
last_offset = row[0] if row else -1
118+
119+
# Idempotency: skip if this offset was already processed.
120+
# This prevents double-counting after a seek-back replays
121+
# messages that were already committed.
122+
if offset <= last_offset:
123+
conn.rollback()
124+
print(f"[skip] offset={offset} already processed (last={last_offset})")
125+
return jsonify(
126+
{
127+
"status": "skipped",
128+
"reason": "already processed",
129+
"offset": offset,
130+
"last_offset": last_offset,
131+
}
132+
)
133+
134+
# Atomically add number to running sum and advance the offset
135+
cur.execute(
136+
"""
137+
UPDATE running_sum
138+
SET total = total + %s,
139+
last_offset = %s,
140+
message_count = message_count + 1
141+
WHERE id = 1
142+
""",
143+
(number, offset),
144+
)
145+
146+
# --- Fault injection ------------------------------------------------
147+
# Simulate a crash between UPDATE and COMMIT.
148+
if FAIL_PROBABILITY > 0 and random.random() < FAIL_PROBABILITY:
149+
raise Exception(
150+
f"Simulated DB failure at offset {offset} "
151+
f"(FAIL_PROBABILITY={FAIL_PROBABILITY})"
152+
)
153+
# --------------------------------------------------------------------
154+
155+
conn.commit()
156+
157+
# Read back the new state for the response
158+
cur.execute(
159+
"SELECT total, last_offset, message_count "
160+
"FROM running_sum WHERE id = 1"
161+
)
162+
total, last_off, count = cur.fetchone()
163+
164+
print(f"[ok] offset={offset} number={number} sum={total} count={count}")
165+
return jsonify(
166+
{
167+
"status": "ok",
168+
"offset": offset,
169+
"number_added": number,
170+
"running_sum": total,
171+
"message_count": count,
172+
}
173+
)
174+
175+
except Exception as e:
176+
print(f"[error] offset={offset} error={e}")
177+
if conn:
178+
try:
179+
conn.rollback()
180+
except Exception:
181+
pass
182+
183+
# Tell OL's Kafka consumer to seek back to this offset and retry.
184+
# The consumer's LRU cache will serve the replay without re-fetching
185+
# from Kafka, and the idempotency check above prevents double-counting
186+
# for any offsets that were already committed before the failure.
187+
resp = make_response(
188+
jsonify({"status": "error", "offset": offset, "error": str(e)}), 500
189+
)
190+
resp.headers["X-Kafka-Seek-Offset"] = str(offset)
191+
return resp
192+
193+
finally:
194+
if conn:
195+
try:
196+
conn.close()
197+
except Exception:
198+
pass

‎examples/kafka-db-sum/instructions.md‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,106 @@
1+
# kafka-db-sum: Testing Instructions
2+
3+
## Prerequisites
4+
5+
- OpenLambda built (`make ol imgs/ol-min`)
6+
- Docker installed
7+
8+
## 1. Start PostgreSQL
9+
10+
```bash
11+
docker run -d --name ol-pg \
12+
--network host \
13+
-e POSTGRES_USER=ol \
14+
-e POSTGRES_PASSWORD=ol \
15+
-e POSTGRES_DB=ol_demo \
16+
postgres:16
17+
```
18+
19+
## 2. Start Kafka
20+
21+
```bash
22+
docker run -d --name kafka \
23+
-p 9092:9092 \
24+
apache/kafka:latest
25+
```
26+
27+
## 3. Create the `numbers` topic
28+
29+
```bash
30+
docker exec kafka /opt/kafka/bin/kafka-topics.sh --create \
31+
--topic numbers \
32+
--bootstrap-server localhost:9092
33+
```
34+
35+
## 4. Initialize and start the OL worker
36+
37+
From the repository root:
38+
39+
```bash
40+
sudo -A ./ol worker init -p ../default-ol -i ol-min
41+
sudo -A ./ol worker up -p ../default-ol
42+
```
43+
44+
Run `worker up` in a separate terminal, or add `-d` for detached mode.
45+
The worker listens on `localhost:5000` by default.
46+
47+
## 5. Install the lambda
48+
49+
From the repository root:
50+
51+
```bash
52+
./ol admin install examples/kafka-db-sum/
53+
```
54+
55+
## 6. Register the Kafka consumer
56+
57+
A standalone worker does not auto-register Kafka triggers on upload.
58+
Register manually:
59+
60+
```bash
61+
curl -X POST localhost:5000/kafka/register/kafka-db-sum
62+
```
63+
64+
## 7. Send test messages
65+
66+
Python producer script (requires `pip install kafka-python`):
67+
68+
```bash
69+
python examples/kafka-db-sum/produce.py 100
70+
```
71+
72+
## 8. Check results
73+
74+
```bash
75+
curl localhost:5000/run/kafka-db-sum/
76+
```
77+
78+
Expected output (sum of 1..100 = 5050):
79+
80+
```json
81+
{ "last_offset": 99, "message_count": 100, "running_sum": 5050 }
82+
```
83+
84+
## 9. Reset and re-run
85+
86+
```bash
87+
curl -X POST localhost:5000/run/kafka-db-sum/reset
88+
```
89+
90+
Then send a fresh batch (step 7) and verify again.
91+
92+
## Configuration
93+
94+
In `ol.yaml`:
95+
96+
| Variable | Default | Description |
97+
| ------------------ | ------------------------------------------- | ------------------------------------------------------------------- |
98+
| `DATABASE_URL` | `postgresql://ol:ol@127.0.0.1:5432/ol_demo` | PostgreSQL connection string |
99+
| `FAIL_PROBABILITY` | `0` | Chance (0.0-1.0) of simulated failure. Use `0.3` to test seek-back. |
100+
101+
## Cleanup
102+
103+
```bash
104+
sudo -A ./ol worker down -p default-ol
105+
docker rm -f kafka ol-pg
106+
```

‎examples/kafka-db-sum/ol.yaml‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
triggers:
2+
http:
3+
- method: GET
4+
- method: POST
5+
kafka:
6+
- bootstrap_servers:
7+
- "localhost:9092"
8+
topics:
9+
- "numbers"
10+
auto_offset_reset: "earliest"
11+
12+
environment:
13+
DATABASE_URL: "postgresql://ol:ol@127.0.0.1:5432/ol_demo"
14+
# Probability (0.0-1.0) of simulated DB failure between UPDATE and COMMIT.
15+
# Set to "0" for normal operation. Try "0.3" to see seek-back recovery in action.
16+
FAIL_PROBABILITY: "0.3"

‎examples/kafka-db-sum/produce.py‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
#!/usr/bin/env python3
2+
"""
3+
Send numbered messages to the 'numbers' Kafka topic.
4+
5+
Usage:
6+
python produce.py # send numbers 1..10
7+
python produce.py 100 # send numbers 1..100
8+
python produce.py 50 0.5 # send 1..50 with 0.5s delay between each
9+
"""
10+
11+
import json
12+
import sys
13+
import time
14+
15+
from kafka import KafkaProducer
16+
17+
BROKER = "localhost:9092"
18+
TOPIC = "numbers"
19+
20+
21+
def main():
22+
count = int(sys.argv[1]) if len(sys.argv) > 1 else 10
23+
delay = float(sys.argv[2]) if len(sys.argv) > 2 else 0.1
24+
25+
producer = KafkaProducer(
26+
bootstrap_servers=BROKER,
27+
value_serializer=lambda v: json.dumps(v).encode("utf-8"),
28+
)
29+
30+
expected_sum = 0
31+
print(f"Sending numbers 1..{count} to topic '{TOPIC}'...")
32+
for i in range(1, count + 1):
33+
producer.send(TOPIC, {"number": i})
34+
expected_sum += i
35+
print(f" sent {i}")
36+
if delay:
37+
time.sleep(delay)
38+
39+
producer.flush()
40+
producer.close()
41+
print(f"\nDone. Expected sum = {expected_sum}")
42+
43+
44+
if __name__ == "__main__":
45+
main()

‎examples/kafka-db-sum/requirements.in‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
flask==2.3.2
2+
werkzeug==3.0.3
3+
psycopg2-binary==2.9.9

‎examples/kafka-db-sum/requirements.txt‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
#
2+
# This file is autogenerated by pip-compile with Python 3.13
3+
# by the following command:
4+
#
5+
# pip-compile requirements.in
6+
#
7+
blinker==1.6.2
8+
# via flask
9+
click==8.1.7
10+
# via flask
11+
flask==2.3.2
12+
# via -r requirements.in
13+
itsdangerous==2.1.2
14+
# via flask
15+
jinja2==3.1.4
16+
# via flask
17+
markupsafe==2.1.3
18+
# via
19+
# jinja2
20+
# werkzeug
21+
psycopg2-binary==2.9.9
22+
# via -r requirements.in
23+
werkzeug==3.0.3
24+
# via
25+
# -r requirements.in
26+
# flask

0 commit comments

Comments
 (0)