1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67

use crate::protocol::EngineProtocol;
use query_structure::PrismaValue;

#[derive(Debug)]
struct RequestContext {
    request_now: PrismaValue,
    engine_protocol: EngineProtocol,
}

tokio::task_local! {
    static REQUEST_CONTEXT: RequestContext;
}

/// A timestamp that should be the `NOW()` value for the whole duration of a request. So all
/// `@default(now())` and `@updatedAt` should use it.
///
/// That panics if REQUEST_CONTEXT has not been set with with_request_context().
///
/// If we had a query context we carry for the entire lifetime of the query, it would belong there.
pub(crate) fn get_request_now() -> PrismaValue {
    REQUEST_CONTEXT.try_with(|rc| rc.request_now.clone()).unwrap_or_else(|_|
            // Task local might not be set in some cases.
            // At the moment of writing, this happens only in query validation test suite.
            // In that case, we want to fall back to realtime value. On the other hand, if task local is
            // set, we want to use it, even if we are not running inside of tokio runtime (for example, 
            // in WASM case)
            //
            // Eventually, this will go away when we have a plain query context reference we pass around.
            PrismaValue::DateTime(chrono::Utc::now().into()))
}

/// The engine protocol used for the whole duration of a request.
/// Use with caution to avoid creating implicit and unnecessary dependencies.
///
/// That panics if REQUEST_CONTEXT has not been set with with_request_context().
///
/// If we had a query context we carry for the entire lifetime of the query, it would belong there.
pub(crate) fn get_engine_protocol() -> EngineProtocol {
    REQUEST_CONTEXT.with(|rc| rc.engine_protocol)
}

/// Execute a future with the current "now" timestamp that can be retrieved through
/// `get_request_now()`, initializing it if necessary.
pub(crate) async fn with_request_context<F, R>(engine_protocol: EngineProtocol, fut: F) -> R
where
    F: std::future::Future<Output = R>,
{
    use chrono::{Duration, DurationRound};

    let is_set = REQUEST_CONTEXT.try_with(|_| async {}).is_ok();

    if is_set {
        fut.await
    } else {
        let timestamp_precision = Duration::milliseconds(1);
        // We round because in create operations, we select after creation and we will fail to
        // select back what we inserted if the timestamp we have is higher precision than the one
        // the database persisted.
        let dt = chrono::Utc::now().duration_round(timestamp_precision).unwrap();
        let ctx = RequestContext {
            request_now: PrismaValue::DateTime(dt.into()),
            engine_protocol,
        };

        REQUEST_CONTEXT.scope(ctx, fut).await
    }
}