Skip to main content

netget/client/dynamodb/
actions.rs

1//! DynamoDB client protocol actions implementation
2
3use crate::llm::actions::{
4    client_trait::{Client, ClientActionResult},
5    protocol_trait::Protocol,
6    ActionDefinition, Parameter, ParameterDefinition,
7};
8use crate::protocol::EventType;
9use crate::state::app_state::AppState;
10use anyhow::{Context, Result};
11use serde_json::json;
12use std::sync::LazyLock;
13
14/// DynamoDB client connected event
15pub static DYNAMODB_CLIENT_CONNECTED_EVENT: LazyLock<EventType> = LazyLock::new(|| {
16    EventType::new(
17        "dynamodb_connected",
18        "DynamoDB client initialized and ready to send requests",
19        json!({"type": "put_item", "table_name": "Users", "item": {"id": {"S": "user456"}, "name": {"S": "Bob"}}}),
20    )
21    .with_parameters(vec![
22        Parameter {
23            name: "region".to_string(),
24            type_hint: "string".to_string(),
25            description: "AWS region for DynamoDB".to_string(),
26            required: true,
27        },
28        Parameter {
29            name: "endpoint".to_string(),
30            type_hint: "string".to_string(),
31            description: "DynamoDB endpoint URL".to_string(),
32            required: false,
33        },
34    ])
35});
36
37/// DynamoDB client response received event
38pub static DYNAMODB_CLIENT_RESPONSE_RECEIVED_EVENT: LazyLock<EventType> = LazyLock::new(|| {
39    EventType::new(
40        "dynamodb_response_received",
41        "DynamoDB response received from server",
42        json!({"type": "query", "table_name": "Users", "key_condition_expression": "id = :id"}),
43    )
44    .with_parameters(vec![
45        Parameter {
46            name: "operation".to_string(),
47            type_hint: "string".to_string(),
48            description: "DynamoDB operation performed".to_string(),
49            required: true,
50        },
51        Parameter {
52            name: "success".to_string(),
53            type_hint: "boolean".to_string(),
54            description: "Whether the operation succeeded".to_string(),
55            required: true,
56        },
57        Parameter {
58            name: "data".to_string(),
59            type_hint: "object".to_string(),
60            description: "Response data from DynamoDB".to_string(),
61            required: false,
62        },
63        Parameter {
64            name: "error".to_string(),
65            type_hint: "string".to_string(),
66            description: "Error message if operation failed".to_string(),
67            required: false,
68        },
69    ])
70});
71
72/// DynamoDB client protocol action handler
73pub struct DynamoDbClientProtocol;
74
75impl DynamoDbClientProtocol {
76    pub fn new() -> Self {
77        Self
78    }
79}
80
81// Implement Protocol trait (common functionality)
82impl Protocol for DynamoDbClientProtocol {
83    fn get_startup_parameters(&self) -> Vec<ParameterDefinition> {
84        vec![
85                ParameterDefinition {
86                    name: "region".to_string(),
87                    description: "AWS region (e.g., us-east-1). Defaults to us-east-1".to_string(),
88                    type_hint: "string".to_string(),
89                    required: false,
90                    example: json!("us-east-1"),
91                    default: None,
92                },
93                ParameterDefinition {
94                    name: "endpoint_url".to_string(),
95                    description: "Custom DynamoDB endpoint URL (for local testing with DynamoDB Local or LocalStack)".to_string(),
96                    type_hint: "string".to_string(),
97                    required: false,
98                    example: json!("http://localhost:8000"),
99                    default: None,
100                },
101                ParameterDefinition {
102                    name: "access_key_id".to_string(),
103                    description: "AWS access key ID (defaults to environment variable)".to_string(),
104                    type_hint: "string".to_string(),
105                    required: false,
106                    example: json!("AKIAIOSFODNN7EXAMPLE"),
107                    default: None,
108                },
109                ParameterDefinition {
110                    name: "secret_access_key".to_string(),
111                    description: "AWS secret access key (defaults to environment variable)".to_string(),
112                    type_hint: "string".to_string(),
113                    required: false,
114                    example: json!("wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"),
115                    default: None,
116                },
117            ]
118    }
119    fn get_async_actions(&self, _state: &AppState) -> Vec<ActionDefinition> {
120        vec![
121            ActionDefinition {
122                name: "put_item".to_string(),
123                description: "Put an item into a DynamoDB table".to_string(),
124                parameters: vec![
125                    Parameter {
126                        name: "table_name".to_string(),
127                        type_hint: "string".to_string(),
128                        description: "Name of the DynamoDB table".to_string(),
129                        required: true,
130                    },
131                    Parameter {
132                        name: "item".to_string(),
133                        type_hint: "object".to_string(),
134                        description: "Item to put (map of attribute names to values with types)"
135                            .to_string(),
136                        required: true,
137                    },
138                ],
139                example: json!({
140                    "type": "put_item",
141                    "table_name": "Users",
142                    "item": {
143                        "id": {"S": "user123"},
144                        "name": {"S": "Alice"},
145                        "age": {"N": "30"}
146                    }
147                }),
148                log_template: None,
149            },
150            ActionDefinition {
151                name: "get_item".to_string(),
152                description: "Get an item from a DynamoDB table by primary key".to_string(),
153                parameters: vec![
154                    Parameter {
155                        name: "table_name".to_string(),
156                        type_hint: "string".to_string(),
157                        description: "Name of the DynamoDB table".to_string(),
158                        required: true,
159                    },
160                    Parameter {
161                        name: "key".to_string(),
162                        type_hint: "object".to_string(),
163                        description:
164                            "Primary key of the item (map of attribute names to values with types)"
165                                .to_string(),
166                        required: true,
167                    },
168                ],
169                example: json!({
170                    "type": "get_item",
171                    "table_name": "Users",
172                    "key": {
173                        "id": {"S": "user123"}
174                    }
175                }),
176                log_template: None,
177            },
178            ActionDefinition {
179                name: "query".to_string(),
180                description: "Query items from a DynamoDB table".to_string(),
181                parameters: vec![
182                    Parameter {
183                        name: "table_name".to_string(),
184                        type_hint: "string".to_string(),
185                        description: "Name of the DynamoDB table".to_string(),
186                        required: true,
187                    },
188                    Parameter {
189                        name: "key_condition_expression".to_string(),
190                        type_hint: "string".to_string(),
191                        description: "Query condition expression".to_string(),
192                        required: true,
193                    },
194                    Parameter {
195                        name: "expression_attribute_values".to_string(),
196                        type_hint: "object".to_string(),
197                        description: "Values for the expression".to_string(),
198                        required: false,
199                    },
200                ],
201                example: json!({
202                    "type": "query",
203                    "table_name": "Users",
204                    "key_condition_expression": "id = :id",
205                    "expression_attribute_values": {
206                        ":id": {"S": "user123"}
207                    }
208                }),
209                log_template: None,
210            },
211            ActionDefinition {
212                name: "scan".to_string(),
213                description: "Scan all items in a DynamoDB table".to_string(),
214                parameters: vec![
215                    Parameter {
216                        name: "table_name".to_string(),
217                        type_hint: "string".to_string(),
218                        description: "Name of the DynamoDB table".to_string(),
219                        required: true,
220                    },
221                    Parameter {
222                        name: "filter_expression".to_string(),
223                        type_hint: "string".to_string(),
224                        description: "Optional filter expression".to_string(),
225                        required: false,
226                    },
227                    Parameter {
228                        name: "expression_attribute_values".to_string(),
229                        type_hint: "object".to_string(),
230                        description: "Values for the filter expression".to_string(),
231                        required: false,
232                    },
233                ],
234                example: json!({
235                    "type": "scan",
236                    "table_name": "Users",
237                    "filter_expression": "age > :min_age",
238                    "expression_attribute_values": {
239                        ":min_age": {"N": "21"}
240                    }
241                }),
242                log_template: None,
243            },
244            ActionDefinition {
245                name: "update_item".to_string(),
246                description: "Update an item in a DynamoDB table".to_string(),
247                parameters: vec![
248                    Parameter {
249                        name: "table_name".to_string(),
250                        type_hint: "string".to_string(),
251                        description: "Name of the DynamoDB table".to_string(),
252                        required: true,
253                    },
254                    Parameter {
255                        name: "key".to_string(),
256                        type_hint: "object".to_string(),
257                        description: "Primary key of the item to update".to_string(),
258                        required: true,
259                    },
260                    Parameter {
261                        name: "update_expression".to_string(),
262                        type_hint: "string".to_string(),
263                        description: "Update expression".to_string(),
264                        required: true,
265                    },
266                    Parameter {
267                        name: "expression_attribute_values".to_string(),
268                        type_hint: "object".to_string(),
269                        description: "Values for the update expression".to_string(),
270                        required: false,
271                    },
272                ],
273                example: json!({
274                    "type": "update_item",
275                    "table_name": "Users",
276                    "key": {
277                        "id": {"S": "user123"}
278                    },
279                    "update_expression": "SET age = :age",
280                    "expression_attribute_values": {
281                        ":age": {"N": "31"}
282                    }
283                }),
284                log_template: None,
285            },
286            ActionDefinition {
287                name: "delete_item".to_string(),
288                description: "Delete an item from a DynamoDB table".to_string(),
289                parameters: vec![
290                    Parameter {
291                        name: "table_name".to_string(),
292                        type_hint: "string".to_string(),
293                        description: "Name of the DynamoDB table".to_string(),
294                        required: true,
295                    },
296                    Parameter {
297                        name: "key".to_string(),
298                        type_hint: "object".to_string(),
299                        description: "Primary key of the item to delete".to_string(),
300                        required: true,
301                    },
302                ],
303                example: json!({
304                    "type": "delete_item",
305                    "table_name": "Users",
306                    "key": {
307                        "id": {"S": "user123"}
308                    }
309                }),
310                log_template: None,
311            },
312            ActionDefinition {
313                name: "disconnect".to_string(),
314                description: "Disconnect from DynamoDB".to_string(),
315                parameters: vec![],
316                example: json!({
317                    "type": "disconnect"
318                }),
319                log_template: None,
320            },
321        ]
322    }
323    fn get_sync_actions(&self) -> Vec<ActionDefinition> {
324        vec![
325            ActionDefinition {
326                name: "put_item".to_string(),
327                description: "Put another item in response to received data".to_string(),
328                parameters: vec![
329                    Parameter {
330                        name: "table_name".to_string(),
331                        type_hint: "string".to_string(),
332                        description: "Name of the DynamoDB table".to_string(),
333                        required: true,
334                    },
335                    Parameter {
336                        name: "item".to_string(),
337                        type_hint: "object".to_string(),
338                        description: "Item to put".to_string(),
339                        required: true,
340                    },
341                ],
342                example: json!({
343                    "type": "put_item",
344                    "table_name": "Users",
345                    "item": {
346                        "id": {"S": "user456"},
347                        "name": {"S": "Bob"}
348                    }
349                }),
350                log_template: None,
351            },
352            ActionDefinition {
353                name: "query".to_string(),
354                description: "Query items in response to received data".to_string(),
355                parameters: vec![
356                    Parameter {
357                        name: "table_name".to_string(),
358                        type_hint: "string".to_string(),
359                        description: "Name of the DynamoDB table".to_string(),
360                        required: true,
361                    },
362                    Parameter {
363                        name: "key_condition_expression".to_string(),
364                        type_hint: "string".to_string(),
365                        description: "Query condition expression".to_string(),
366                        required: true,
367                    },
368                ],
369                example: json!({
370                    "type": "query",
371                    "table_name": "Users",
372                    "key_condition_expression": "id = :id"
373                }),
374                log_template: None,
375            },
376        ]
377    }
378    fn protocol_name(&self) -> &'static str {
379        "DynamoDB"
380    }
381    fn get_event_types(&self) -> Vec<EventType> {
382        vec![
383            EventType::new(
384                "dynamodb_connected",
385                "Triggered when DynamoDB client is initialized",
386                json!({"type": "put_item", "table_name": "Users", "item": {"id": {"S": "user456"}, "name": {"S": "Bob"}}}),
387            ),
388            EventType::new(
389                "dynamodb_response_received",
390                "Triggered when DynamoDB client receives a response",
391                json!({"type": "query", "table_name": "Users", "key_condition_expression": "id = :id"}),
392            ),
393        ]
394    }
395    fn stack_name(&self) -> &'static str {
396        "ETH>IP>TCP>TLS>HTTP>DynamoDB"
397    }
398    fn keywords(&self) -> Vec<&'static str> {
399        vec!["dynamodb", "dynamo", "aws dynamodb", "connect to dynamodb"]
400    }
401    fn metadata(&self) -> crate::protocol::metadata::ProtocolMetadataV2 {
402        use crate::protocol::metadata::{DevelopmentState, ProtocolMetadataV2};
403
404        ProtocolMetadataV2::builder()
405                .state(DevelopmentState::Experimental)
406                .implementation("AWS SDK for DynamoDB client library")
407                .llm_control("Full control over DynamoDB operations (PutItem, GetItem, Query, Scan, UpdateItem, DeleteItem)")
408                .e2e_testing("DynamoDB Local or LocalStack")
409                .build()
410    }
411    fn description(&self) -> &'static str {
412        "DynamoDB client for interacting with AWS DynamoDB or local instances"
413    }
414    fn example_prompt(&self) -> &'static str {
415        "Connect to DynamoDB and put an item in the Users table"
416    }
417    fn group_name(&self) -> &'static str {
418        "Database"
419    }
420
421    fn get_startup_examples(&self) -> crate::llm::actions::StartupExamples {
422        use crate::llm::actions::StartupExamples;
423        use serde_json::json;
424
425        // A custom endpoint is refused without explicit credentials (`client::aws_support`),
426        // so every example names a dummy pair, which is what DynamoDB Local accepts.
427        StartupExamples::new(
428            // LLM mode: LLM controls DynamoDB operations
429            json!({
430                "type": "open_client",
431                "remote_addr": "localhost:8000",
432                "base_stack": "dynamodb",
433                "startup_params": {
434                    "access_key_id": "local",
435                    "secret_access_key": "local"
436                },
437                "instruction": "Scan the Users table and report all items with age greater than 21"
438            }),
439            // Script mode: Code-based deterministic responses
440            json!({
441                "type": "open_client",
442                "remote_addr": "localhost:8000",
443                "base_stack": "dynamodb",
444                "startup_params": {
445                    "access_key_id": "local",
446                    "secret_access_key": "local"
447                },
448                "event_handlers": [{
449                    "event_pattern": "dynamodb_response_received",
450                    "handler": {
451                        "type": "script",
452                        "language": "python",
453                        "code": "<dynamodb_client_handler>"
454                    }
455                }]
456            }),
457            // Static mode: Fixed DynamoDB scan on connect
458            json!({
459                "type": "open_client",
460                "remote_addr": "localhost:8000",
461                "base_stack": "dynamodb",
462                "startup_params": {
463                    "access_key_id": "local",
464                    "secret_access_key": "local"
465                },
466                "event_handlers": [
467                    {
468                        "event_pattern": "dynamodb_connected",
469                        "handler": {
470                            "type": "static",
471                            "actions": [{
472                                "type": "scan",
473                                "table_name": "Users"
474                            }]
475                        }
476                    },
477                    {
478                        "event_pattern": "dynamodb_response_received",
479                        "handler": {
480                            "type": "static",
481                            "actions": [{
482                                "type": "disconnect"
483                            }]
484                        }
485                    }
486                ]
487            }),
488        )
489    }
490}
491
492// Implement Client trait (client-specific functionality)
493impl Client for DynamoDbClientProtocol {
494    fn connect(
495        &self,
496        ctx: crate::protocol::ConnectContext,
497    ) -> std::pin::Pin<
498        Box<dyn std::future::Future<Output = anyhow::Result<std::net::SocketAddr>> + Send>,
499    > {
500        Box::pin(async move {
501            use crate::client::dynamodb::DynamoDbClient;
502            DynamoDbClient::connect_with_llm_actions(
503                ctx.remote_addr,
504                ctx.llm_client,
505                ctx.state,
506                ctx.status_tx,
507                ctx.client_id,
508                ctx.startup_params,
509            )
510            .await
511        })
512    }
513    fn execute_action(&self, action: serde_json::Value) -> Result<ClientActionResult> {
514        let action_type = action
515            .get("type")
516            .and_then(|v| v.as_str())
517            .context("Missing 'type' field in action")?;
518
519        match action_type {
520            "put_item" => {
521                let table_name = action
522                    .get("table_name")
523                    .and_then(|v| v.as_str())
524                    .context("Missing 'table_name' field")?
525                    .to_string();
526
527                let item = action
528                    .get("item")
529                    .and_then(|v| v.as_object())
530                    .context("Missing 'item' field")?
531                    .clone();
532
533                Ok(ClientActionResult::Custom {
534                    name: "put_item".to_string(),
535                    data: json!({
536                        "table_name": table_name,
537                        "item": item,
538                    }),
539                })
540            }
541            "get_item" => {
542                let table_name = action
543                    .get("table_name")
544                    .and_then(|v| v.as_str())
545                    .context("Missing 'table_name' field")?
546                    .to_string();
547
548                let key = action
549                    .get("key")
550                    .and_then(|v| v.as_object())
551                    .context("Missing 'key' field")?
552                    .clone();
553
554                Ok(ClientActionResult::Custom {
555                    name: "get_item".to_string(),
556                    data: json!({
557                        "table_name": table_name,
558                        "key": key,
559                    }),
560                })
561            }
562            "query" => {
563                let table_name = action
564                    .get("table_name")
565                    .and_then(|v| v.as_str())
566                    .context("Missing 'table_name' field")?
567                    .to_string();
568
569                let key_condition_expression = action
570                    .get("key_condition_expression")
571                    .and_then(|v| v.as_str())
572                    .context("Missing 'key_condition_expression' field")?
573                    .to_string();
574
575                let expression_attribute_values = action
576                    .get("expression_attribute_values")
577                    .and_then(|v| v.as_object())
578                    .cloned();
579
580                Ok(ClientActionResult::Custom {
581                    name: "query".to_string(),
582                    data: json!({
583                        "table_name": table_name,
584                        "key_condition_expression": key_condition_expression,
585                        "expression_attribute_values": expression_attribute_values,
586                    }),
587                })
588            }
589            "scan" => {
590                let table_name = action
591                    .get("table_name")
592                    .and_then(|v| v.as_str())
593                    .context("Missing 'table_name' field")?
594                    .to_string();
595
596                let filter_expression = action
597                    .get("filter_expression")
598                    .and_then(|v| v.as_str())
599                    .map(|s| s.to_string());
600
601                let expression_attribute_values = action
602                    .get("expression_attribute_values")
603                    .and_then(|v| v.as_object())
604                    .cloned();
605
606                Ok(ClientActionResult::Custom {
607                    name: "scan".to_string(),
608                    data: json!({
609                        "table_name": table_name,
610                        "filter_expression": filter_expression,
611                        "expression_attribute_values": expression_attribute_values,
612                    }),
613                })
614            }
615            "update_item" => {
616                let table_name = action
617                    .get("table_name")
618                    .and_then(|v| v.as_str())
619                    .context("Missing 'table_name' field")?
620                    .to_string();
621
622                let key = action
623                    .get("key")
624                    .and_then(|v| v.as_object())
625                    .context("Missing 'key' field")?
626                    .clone();
627
628                let update_expression = action
629                    .get("update_expression")
630                    .and_then(|v| v.as_str())
631                    .context("Missing 'update_expression' field")?
632                    .to_string();
633
634                let expression_attribute_values = action
635                    .get("expression_attribute_values")
636                    .and_then(|v| v.as_object())
637                    .cloned();
638
639                Ok(ClientActionResult::Custom {
640                    name: "update_item".to_string(),
641                    data: json!({
642                        "table_name": table_name,
643                        "key": key,
644                        "update_expression": update_expression,
645                        "expression_attribute_values": expression_attribute_values,
646                    }),
647                })
648            }
649            "delete_item" => {
650                let table_name = action
651                    .get("table_name")
652                    .and_then(|v| v.as_str())
653                    .context("Missing 'table_name' field")?
654                    .to_string();
655
656                let key = action
657                    .get("key")
658                    .and_then(|v| v.as_object())
659                    .context("Missing 'key' field")?
660                    .clone();
661
662                Ok(ClientActionResult::Custom {
663                    name: "delete_item".to_string(),
664                    data: json!({
665                        "table_name": table_name,
666                        "key": key,
667                    }),
668                })
669            }
670            "disconnect" => Ok(ClientActionResult::Disconnect),
671            _ => Err(anyhow::anyhow!(
672                "Unknown DynamoDB client action: {}",
673                action_type
674            )),
675        }
676    }
677}