Skip to main content

tank_tests/
trade.rs

1#![allow(unused_imports)]
2use rust_decimal::Decimal;
3use std::{collections::BTreeMap, pin::pin, str::FromStr, sync::LazyLock};
4use tank::{
5    AsValue, Driver, DynQuery, Entity, Executor, FixedDecimal, Passive, Query, QueryBuilder,
6    QueryResult, RawQuery, RowsAffected, SqlWriter, Value,
7    stream::{StreamExt, TryStreamExt},
8};
9use time::macros::datetime;
10use tokio::sync::Mutex;
11use uuid::Uuid;
12
13static MUTEX: LazyLock<Mutex<()>> = LazyLock::new(|| Mutex::new(()));
14
15#[derive(Entity, Debug, PartialEq)]
16#[tank(schema = "trading", name = "trade_execution", primary_key = ("trade_id", "execution_time"))]
17pub struct Trade {
18    #[tank(name = "trade_id")]
19    pub trade: u64,
20    #[tank(name = "order_id", default = Uuid::from_str("241d362d-797e-4769-b3f6-412440c8cf68").unwrap().as_value())]
21    pub order: Uuid,
22    /// Ticker symbol
23    pub symbol: String,
24    #[cfg(not(feature = "disable-arrays"))]
25    pub isin: [char; 12],
26    pub price: FixedDecimal<18, 4>,
27    pub quantity: u32,
28    pub execution_time: Passive<time::PrimitiveDateTime>,
29    pub currency: Option<String>,
30    pub is_internalized: bool,
31    /// Exchange
32    pub venue: Option<String>,
33    #[cfg(not(feature = "disable-lists"))]
34    pub child_trade_ids: Option<Vec<i64>>,
35    pub metadata: Option<Box<[u8]>>,
36    #[cfg(not(feature = "disable-maps"))]
37    pub tags: Option<BTreeMap<String, String>>,
38}
39
40pub async fn trade_simple(executor: &mut impl Executor) {
41    let _lock = MUTEX.lock().await;
42
43    // Setup
44    Trade::drop_table(executor, true, false)
45        .await
46        .expect("Failed to drop Trade table");
47    Trade::create_table(executor, false, true)
48        .await
49        .expect("Failed to create Trade table");
50
51    // Trade object
52    let trade = Trade {
53        trade: 46923,
54        order: Uuid::parse_str("550e8400-e29b-41d4-a716-446655440000").unwrap(),
55        symbol: "RIVN".to_string(),
56        #[cfg(not(feature = "disable-arrays"))]
57        isin: std::array::from_fn(|i| "US76954A1034".chars().nth(i).unwrap()),
58        price: Decimal::new(1226, 2).into(), // 12.26
59        quantity: 500,
60        execution_time: datetime!(2025-06-07 14:32:00).into(),
61        currency: Some("USD".into()),
62        is_internalized: true,
63        venue: Some("NASDAQ".into()),
64        #[cfg(not(feature = "disable-lists"))]
65        child_trade_ids: vec![36209, 85320].into(),
66        metadata: b"Metadata Bytes".to_vec().into_boxed_slice().into(),
67        #[cfg(not(feature = "disable-maps"))]
68        tags: BTreeMap::from_iter([
69            ("source".into(), "internal".into()),
70            ("strategy".into(), "scalping".into()),
71        ])
72        .into(),
73    };
74
75    // Expect to find no trades
76    let result = Trade::find_one(executor, trade.primary_key_expr())
77        .await
78        .expect("Failed to find trade by primary key");
79    assert!(result.is_none(), "Expected no trades at this time");
80    assert_eq!(
81        Trade::find_many(executor, true, None)
82            .map_err(|e| panic!("{e:#}"))
83            .count()
84            .await,
85        0
86    );
87
88    // Save a trade
89    trade.save(executor).await.expect("Failed to save trade");
90
91    // Expect to find the only trade
92    let result = Trade::find_one(executor, trade.primary_key_expr())
93        .await
94        .expect("Failed to find trade");
95    assert!(
96        result.is_some(),
97        "Expected Trade::find_one to return some result",
98    );
99    let result = result.unwrap();
100    assert_eq!(result.trade, 46923);
101    assert_eq!(
102        result.order,
103        Uuid::from_str("550e8400-e29b-41d4-a716-446655440000").unwrap()
104    );
105    assert_eq!(result.symbol, "RIVN");
106    #[cfg(not(feature = "disable-arrays"))]
107    assert_eq!(
108        result
109            .isin
110            .iter()
111            .map(|v| v.to_string())
112            .collect::<Vec<_>>()
113            .join(""),
114        "US76954A1034"
115    );
116    assert_eq!(result.price, Decimal::new(1226, 2).into());
117    assert_eq!(result.quantity, 500);
118    assert_eq!(
119        result.execution_time,
120        Passive::Set(datetime!(2025-06-07 14:32:00))
121    );
122    assert_eq!(result.currency, Some("USD".into()));
123    assert_eq!(result.is_internalized, true);
124    assert_eq!(result.venue, Some("NASDAQ".into()));
125    #[cfg(not(feature = "disable-lists"))]
126    assert_eq!(result.child_trade_ids, Some(vec![36209, 85320]));
127    assert_eq!(
128        result.metadata,
129        Some(b"Metadata Bytes".to_vec().into_boxed_slice())
130    );
131    #[cfg(not(feature = "disable-maps"))]
132    let Some(tags) = result.tags else {
133        unreachable!("Tag is expected");
134    };
135    #[cfg(not(feature = "disable-maps"))]
136    assert_eq!(tags.len(), 2);
137    #[cfg(not(feature = "disable-maps"))]
138    assert_eq!(
139        tags,
140        BTreeMap::from_iter([
141            ("source".into(), "internal".into()),
142            ("strategy".into(), "scalping".into())
143        ])
144    );
145
146    assert_eq!(Trade::find_many(executor, true, None).count().await, 1);
147}
148
149pub async fn trade_multiple(executor: &mut impl Executor) {
150    let _lock = MUTEX.lock().await;
151
152    // Setup
153    Trade::drop_table(executor, false, false)
154        .await
155        .expect("Failed to drop Trade table");
156    Trade::create_table(executor, false, true)
157        .await
158        .expect("Failed to create Trade table");
159
160    // Trade objects
161    let trades = vec![
162        Trade {
163            trade: 10001,
164            order: Uuid::parse_str("11111111-1111-1111-1111-111111111111").unwrap(),
165            symbol: "AAPL".to_string(),
166            #[cfg(not(feature = "disable-arrays"))]
167            isin: std::array::from_fn(|i| "US0378331005".chars().nth(i).unwrap()),
168            price: Decimal::new(15000, 2).into(),
169            quantity: 10,
170            execution_time: datetime!(2025-06-01 09:00:00).into(),
171            currency: Some("USD".into()),
172            is_internalized: false,
173            venue: Some("NASDAQ".into()),
174            #[cfg(not(feature = "disable-lists"))]
175            child_trade_ids: Some(vec![101, 102]),
176            metadata: Some(b"First execution".to_vec().into_boxed_slice()),
177            #[cfg(not(feature = "disable-maps"))]
178            tags: Some(BTreeMap::from_iter([
179                ("source".into(), "algo".into()),
180                ("strategy".into(), "momentum".into()),
181            ])),
182        },
183        Trade {
184            trade: 10002,
185            order: Uuid::parse_str("22222222-2222-2222-2222-222222222222").unwrap(),
186            symbol: "GOOG".to_string(),
187            #[cfg(not(feature = "disable-arrays"))]
188            isin: std::array::from_fn(|i| "US02079K3059".chars().nth(i).unwrap()),
189            price: Decimal::new(280000, 3).into(), // 280.000
190            quantity: 5,
191            execution_time: datetime!(2025-06-02 10:15:30).into(),
192            currency: Some("USD".into()),
193            is_internalized: true,
194            venue: Some("NYSE".into()),
195            #[cfg(not(feature = "disable-lists"))]
196            child_trade_ids: None,
197            metadata: Some(b"Second execution".to_vec().into_boxed_slice()),
198            #[cfg(not(feature = "disable-maps"))]
199            tags: Some(BTreeMap::from_iter([
200                ("source".into(), "internal".into()),
201                ("strategy".into(), "mean_reversion".into()),
202            ])),
203        },
204        Trade {
205            trade: 10003,
206            order: Uuid::parse_str("33333333-3333-3333-3333-333333333333").unwrap(),
207            symbol: "MSFT".to_string(),
208            #[cfg(not(feature = "disable-arrays"))]
209            isin: std::array::from_fn(|i| "US5949181045".chars().nth(i).unwrap()),
210            price: Decimal::new(32567, 2).into(), // 325.67
211            quantity: 20,
212            execution_time: datetime!(2025-06-03 11:45:00).into(),
213            currency: Some("USD".into()),
214            is_internalized: false,
215            venue: Some("BATS".into()),
216            #[cfg(not(feature = "disable-lists"))]
217            child_trade_ids: Some(vec![301]),
218            metadata: Some(b"Third execution".to_vec().into_boxed_slice()),
219            #[cfg(not(feature = "disable-maps"))]
220            tags: Some(BTreeMap::from_iter([
221                ("sourcev".into(), "external".into()),
222                ("strategy".into(), "arbitrage".into()),
223            ])),
224        },
225        Trade {
226            trade: 10004,
227            order: Uuid::parse_str("44444444-4444-4444-4444-444444444444").unwrap(),
228            symbol: "TSLA".to_string(),
229            #[cfg(not(feature = "disable-arrays"))]
230            isin: std::array::from_fn(|i| "US88160R1014".chars().nth(i).unwrap()),
231            price: Decimal::new(62000, 2).into(), // 620.00
232            quantity: 15,
233            execution_time: datetime!(2025-06-04 14:00:00).into(),
234            currency: Some("USD".into()),
235            is_internalized: true,
236            venue: Some("CBOE".into()),
237            #[cfg(not(feature = "disable-lists"))]
238            child_trade_ids: None,
239            metadata: None,
240            #[cfg(not(feature = "disable-maps"))]
241            tags: Some(BTreeMap::from_iter([
242                ("source".into(), "manual".into()),
243                ("strategy".into(), "news_event".into()),
244            ])),
245        },
246        Trade {
247            trade: 10005,
248            order: Uuid::parse_str("55555555-5555-5555-5555-555555555555").unwrap(),
249            symbol: "AMZN".to_string(),
250            #[cfg(not(feature = "disable-arrays"))]
251            isin: std::array::from_fn(|i| "US0231351067".chars().nth(i).unwrap()),
252            price: Decimal::new(134899, 3).into(), // 1348.99
253            quantity: 8,
254            execution_time: datetime!(2025-06-05 16:30:00).into(),
255            currency: Some("USD".into()),
256            is_internalized: false,
257            venue: Some("NASDAQ".into()),
258            #[cfg(not(feature = "disable-lists"))]
259            child_trade_ids: Some(vec![501, 502, 503]),
260            metadata: Some(b"Fifth execution".to_vec().into_boxed_slice()),
261            #[cfg(not(feature = "disable-maps"))]
262            tags: Some(BTreeMap::from_iter([
263                ("source".into(), "internal".into()),
264                ("strategy".into(), "scalping".into()),
265            ])),
266        },
267    ];
268
269    // Insert 5 trades
270    let affected = Trade::insert_many(executor, &trades)
271        .await
272        .expect("Coult not insert 5 trade");
273    if let Some(affected) = affected.rows_affected {
274        assert_eq!(affected, 5);
275    }
276
277    // Find 5 trades
278    let data = Trade::find_many(executor, true, None)
279        .try_collect::<Vec<_>>()
280        .await
281        .expect("Failed to query threads");
282    assert_eq!(data.len(), 5, "Expect to find 5 trades");
283
284    // Verify data integrity
285    for (i, expected) in trades.iter().enumerate() {
286        let actual_a = &trades[i];
287        let actual_b = Trade::find_one(executor, expected.primary_key_expr())
288            .await
289            .expect(&format!("Failed to find trade {} by pk", data[i].symbol));
290        let Some(actual_b) = actual_b else {
291            panic!("Trade {} not found", expected.trade);
292        };
293
294        assert_eq!(actual_a.trade, expected.trade);
295        assert_eq!(actual_b.trade, expected.trade);
296
297        assert_eq!(actual_a.order, expected.order);
298        assert_eq!(actual_b.order, expected.order);
299
300        assert_eq!(actual_a.symbol, expected.symbol);
301        assert_eq!(actual_b.symbol, expected.symbol);
302
303        assert_eq!(actual_a.price, expected.price);
304        assert_eq!(actual_b.price, expected.price);
305
306        assert_eq!(actual_a.quantity, expected.quantity);
307        assert_eq!(actual_b.quantity, expected.quantity);
308
309        assert_eq!(actual_a.execution_time, expected.execution_time);
310        assert_eq!(actual_b.execution_time, expected.execution_time);
311
312        assert_eq!(actual_a.currency, expected.currency);
313        assert_eq!(actual_b.currency, expected.currency);
314
315        assert_eq!(actual_a.is_internalized, expected.is_internalized);
316        assert_eq!(actual_b.is_internalized, expected.is_internalized);
317
318        assert_eq!(actual_a.venue, expected.venue);
319        assert_eq!(actual_b.venue, expected.venue);
320
321        #[cfg(not(feature = "disable-lists"))]
322        assert_eq!(actual_a.child_trade_ids, expected.child_trade_ids);
323        #[cfg(not(feature = "disable-lists"))]
324        assert_eq!(actual_b.child_trade_ids, expected.child_trade_ids);
325
326        assert_eq!(actual_a.metadata, expected.metadata);
327        assert_eq!(actual_b.metadata, expected.metadata);
328
329        #[cfg(not(feature = "disable-maps"))]
330        assert_eq!(actual_a.tags, expected.tags);
331        #[cfg(not(feature = "disable-maps"))]
332        assert_eq!(actual_b.tags, expected.tags);
333    }
334
335    // Multiple statements
336    #[cfg(not(feature = "disable-multiple-statements"))]
337    {
338        let writer = executor.driver().sql_writer();
339        let mut query = DynQuery::default();
340        writer.write_delete::<Trade>(&mut query, true);
341        writer.write_insert(
342            &mut query,
343            &[Trade {
344                trade: 10002,
345                order: Uuid::parse_str("895dc048-be92-4a55-afbf-38a60936e844").unwrap(),
346                symbol: "RIVN".to_string(),
347                #[cfg(not(feature = "disable-arrays"))]
348                isin: std::array::from_fn(|i| "US76954A1034".chars().nth(i).unwrap()),
349                price: Decimal::new(1345, 2).into(),
350                quantity: 3200,
351                execution_time: datetime!(2025-06-01 10:15:30).into(),
352                currency: Some("USD".into()),
353                is_internalized: true,
354                venue: Some("NASDAQ".into()),
355                #[cfg(not(feature = "disable-lists"))]
356                child_trade_ids: Some(vec![201]),
357                metadata: Some(
358                    b"desc: \"Crossed with internal liquidity\", id:'\\X696E7465726E616C'"
359                        .to_vec()
360                        .into_boxed_slice(),
361                ),
362                #[cfg(not(feature = "disable-maps"))]
363                tags: Some(BTreeMap::from_iter([
364                    ("source".into(), "internal".into()),
365                    ("strategy".into(), "arbitrage".into()),
366                    ("risk_limit".into(), "high".into()),
367                ])),
368            }],
369            false,
370        );
371        writer.write_select(
372            &mut query,
373            &QueryBuilder::new()
374                .select(Trade::columns())
375                .from(Trade::table())
376                .where_expr(true),
377        );
378        let mut stream = pin!(executor.run(query));
379        let Some(Ok(QueryResult::Affected(RowsAffected { rows_affected, .. }))) =
380            stream.next().await
381        else {
382            panic!("Could not get the result of the first query");
383        };
384        if let Some(rows_affected) = rows_affected {
385            assert_eq!(rows_affected, 5);
386        }
387        let Some(Ok(QueryResult::Affected(RowsAffected { rows_affected, .. }))) =
388            stream.next().await
389        else {
390            panic!("Could not get the result of the first statement");
391        };
392        if let Some(rows_affected) = rows_affected {
393            assert_eq!(rows_affected, 1);
394        }
395        let Some(Ok(QueryResult::Row(row))) = stream.next().await else {
396            panic!("Could not get the result of the second statement");
397        };
398        let trade = Trade::from_row(row).expect("Could not decode the Trade from row");
399        assert_eq!(
400            trade,
401            Trade {
402                trade: 10002,
403                order: Uuid::parse_str("895dc048-be92-4a55-afbf-38a60936e844").unwrap(),
404                symbol: "RIVN".to_string(),
405                #[cfg(not(feature = "disable-arrays"))]
406                isin: std::array::from_fn(|i| "US76954A1034".chars().nth(i).unwrap()),
407                price: Decimal::new(1345, 2).into(),
408                quantity: 3200,
409                execution_time: datetime!(2025-06-01 10:15:30).into(),
410                currency: Some("USD".into()),
411                is_internalized: true,
412                venue: Some("NASDAQ".into()),
413                #[cfg(not(feature = "disable-lists"))]
414                child_trade_ids: Some(vec![201]),
415                metadata: Some(
416                    b"desc: \"Crossed with internal liquidity\", id:'\\X696E7465726E616C'"
417                        .to_vec()
418                        .into_boxed_slice(),
419                ),
420                #[cfg(not(feature = "disable-maps"))]
421                tags: Some(BTreeMap::from_iter([
422                    ("source".into(), "internal".into()),
423                    ("strategy".into(), "arbitrage".into()),
424                    ("risk_limit".into(), "high".into()),
425                ])),
426            }
427        );
428    }
429}