Skip to main content

objectiveai_sdk/cli/command/db/
mod.rs

1//! `db` — direct database access at the CLI level.
2//!
3//! One leaf today:
4//! - `query` — execute arbitrary single-statement read-only SQL
5//!   with a required timeout and an optional per-response token
6//!   budget.
7
8use crate::cli::command::CommandRequest;
9
10pub mod query;
11
12#[derive(clap::Subcommand)]
13pub enum Command {
14    /// Execute an arbitrary single-statement read-only SQL query
15    /// against the CLI's local postgres pool.
16    Query(query::Command),
17}
18
19#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
20#[serde(untagged)]
21#[schemars(rename = "cli.command.db.Request")]
22pub enum Request {
23    #[schemars(title = "Query")]
24    Query(query::Request),
25    #[schemars(title = "QueryRequestSchema")]
26    QueryRequestSchema(query::request_schema::Request),
27    #[schemars(title = "QueryResponseSchema")]
28    QueryResponseSchema(query::response_schema::Request),
29}
30
31// Exempt from json-schema coverage: tier aggregate (see the root
32// `ResponseItem` in command.rs - TS7056).
33#[objectiveai_sdk_macros::json_schema_ignore]
34#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, schemars::JsonSchema)]
35#[schemars(rename = "cli.command.db.ResponseItem")]
36#[serde(untagged)]
37pub enum ResponseItem {
38    #[schemars(title = "Query")]
39    Query(query::Response),
40    #[schemars(title = "QueryRequestSchema")]
41    QueryRequestSchema(query::request_schema::Response),
42    #[schemars(title = "QueryResponseSchema")]
43    QueryResponseSchema(query::response_schema::Response),
44}
45
46#[cfg(feature = "mcp")]
47impl crate::cli::command::CommandResponse for ResponseItem {
48    fn into_mcp(self) -> crate::cli::command::McpResponseItem {
49        match self {
50            ResponseItem::Query(v) => v.into_mcp(),
51            ResponseItem::QueryRequestSchema(v) => v.into_mcp(),
52            ResponseItem::QueryResponseSchema(v) => v.into_mcp(),
53        }
54    }
55}
56
57impl TryFrom<Command> for Request {
58    type Error = crate::cli::command::FromArgsError;
59    fn try_from(command: Command) -> Result<Self, Self::Error> {
60        match command {
61            Command::Query(cmd) => match cmd.schema {
62                None => Ok(Request::Query(query::Request::try_from(cmd.args)?)),
63                Some(query::Schema::RequestSchema(args)) => Ok(
64                    Request::QueryRequestSchema(query::request_schema::Request::try_from(args)?),
65                ),
66                Some(query::Schema::ResponseSchema(args)) => Ok(
67                    Request::QueryResponseSchema(query::response_schema::Request::try_from(args)?),
68                ),
69            },
70        }
71    }
72}
73
74impl CommandRequest for Request {
75    fn into_command(&self) -> Vec<String> {
76        match self {
77            Request::Query(inner) => inner.into_command(),
78            Request::QueryRequestSchema(inner) => inner.into_command(),
79            Request::QueryResponseSchema(inner) => inner.into_command(),
80        }
81    }
82}
83
84#[cfg(feature = "cli-executor")]
85pub async fn execute<E: crate::cli::command::CommandExecutor>(
86    executor: &E,
87    request: Request,
88    agent_arguments: Option<&crate::cli::command::AgentArguments>,
89) -> Result<
90    std::pin::Pin<Box<dyn futures::Stream<Item = Result<ResponseItem, E::Error>> + Send>>,
91    E::Error,
92> {
93    let stream: std::pin::Pin<
94        Box<dyn futures::Stream<Item = Result<ResponseItem, E::Error>> + Send>,
95    > = match request {
96        Request::Query(req) => {
97            let value = query::execute(executor, req, agent_arguments).await?;
98            Box::pin(crate::cli::command::StreamOnce::new(Ok(
99                ResponseItem::Query(value),
100            )))
101        }
102        Request::QueryRequestSchema(req) => {
103            let value = query::request_schema::execute(executor, req, agent_arguments).await?;
104            Box::pin(crate::cli::command::StreamOnce::new(Ok(
105                ResponseItem::QueryRequestSchema(value),
106            )))
107        }
108        Request::QueryResponseSchema(req) => {
109            let value = query::response_schema::execute(executor, req, agent_arguments).await?;
110            Box::pin(crate::cli::command::StreamOnce::new(Ok(
111                ResponseItem::QueryResponseSchema(value),
112            )))
113        }
114    };
115    Ok(stream)
116}
117
118#[cfg(feature = "cli-executor")]
119pub async fn execute_jq<E: crate::cli::command::CommandExecutor>(
120    executor: &E,
121    request: Request,
122    jq: String,
123    agent_arguments: Option<&crate::cli::command::AgentArguments>,
124) -> Result<
125    std::pin::Pin<Box<dyn futures::Stream<Item = Result<serde_json::Value, E::Error>> + Send>>,
126    E::Error,
127> {
128    let stream: std::pin::Pin<
129        Box<dyn futures::Stream<Item = Result<serde_json::Value, E::Error>> + Send>,
130    > = match request {
131        Request::Query(req) => {
132            let value = query::execute_jq(executor, req, jq, agent_arguments).await?;
133            Box::pin(crate::cli::command::StreamOnce::new(Ok(value)))
134        }
135        Request::QueryRequestSchema(req) => {
136            let value =
137                query::request_schema::execute_jq(executor, req, jq, agent_arguments).await?;
138            Box::pin(crate::cli::command::StreamOnce::new(Ok(value)))
139        }
140        Request::QueryResponseSchema(req) => {
141            let value =
142                query::response_schema::execute_jq(executor, req, jq, agent_arguments).await?;
143            Box::pin(crate::cli::command::StreamOnce::new(Ok(value)))
144        }
145    };
146    Ok(stream)
147}