nirv_engine/protocol/
protocol_trait.rs

1use async_trait::async_trait;
2use std::collections::HashMap;
3use tokio::net::TcpStream;
4use crate::utils::{NirvResult, ProtocolError, InternalQuery, QueryResult};
5
6/// Protocol types supported by NIRV Engine
7#[derive(Debug, Clone, PartialEq, Eq, Hash)]
8pub enum ProtocolType {
9    PostgreSQL,
10    MySQL,
11    SQLite,
12    SqlServer,
13}
14
15/// Connection state for protocol adapters
16#[derive(Debug)]
17pub struct Connection {
18    pub stream: TcpStream,
19    pub authenticated: bool,
20    pub database: String,
21    pub parameters: HashMap<String, String>,
22    pub protocol_type: ProtocolType,
23}
24
25impl Connection {
26    pub fn new(stream: TcpStream, protocol_type: ProtocolType) -> Self {
27        Self {
28            stream,
29            authenticated: false,
30            database: String::new(),
31            parameters: HashMap::new(),
32            protocol_type,
33        }
34    }
35}
36
37/// Authentication credentials
38#[derive(Debug, Clone)]
39pub struct Credentials {
40    pub username: String,
41    pub password: Option<String>,
42    pub database: String,
43    pub parameters: HashMap<String, String>,
44}
45
46impl Credentials {
47    pub fn new(username: String, database: String) -> Self {
48        Self {
49            username,
50            password: None,
51            database,
52            parameters: HashMap::new(),
53        }
54    }
55    
56    pub fn with_password(mut self, password: String) -> Self {
57        self.password = Some(password);
58        self
59    }
60    
61    pub fn with_parameter(mut self, key: String, value: String) -> Self {
62        self.parameters.insert(key, value);
63        self
64    }
65}
66
67/// Protocol-specific query representation
68#[derive(Debug, Clone)]
69pub struct ProtocolQuery {
70    pub raw_query: String,
71    pub parameters: Vec<String>,
72    pub protocol_type: ProtocolType,
73}
74
75impl ProtocolQuery {
76    pub fn new(raw_query: String, protocol_type: ProtocolType) -> Self {
77        Self {
78            raw_query,
79            parameters: Vec::new(),
80            protocol_type,
81        }
82    }
83    
84    pub fn with_parameters(mut self, parameters: Vec<String>) -> Self {
85        self.parameters = parameters;
86        self
87    }
88}
89
90/// Protocol-specific response representation
91#[derive(Debug, Clone)]
92pub struct ProtocolResponse {
93    pub result: QueryResult,
94    pub protocol_type: ProtocolType,
95    pub format: ResponseFormat,
96}
97
98/// Response format options
99#[derive(Debug, Clone, PartialEq)]
100pub enum ResponseFormat {
101    Text,
102    Binary,
103}
104
105impl ProtocolResponse {
106    pub fn new(result: QueryResult, protocol_type: ProtocolType) -> Self {
107        Self {
108            result,
109            protocol_type,
110            format: ResponseFormat::Text,
111        }
112    }
113    
114    pub fn with_format(mut self, format: ResponseFormat) -> Self {
115        self.format = format;
116        self
117    }
118}
119
120/// Main trait for database protocol adapters
121#[async_trait]
122pub trait ProtocolAdapter: Send + Sync {
123    /// Accept a new connection and perform initial handshake
124    async fn accept_connection(&self, stream: TcpStream) -> NirvResult<Connection>;
125    
126    /// Authenticate a connection with provided credentials
127    async fn authenticate(&self, conn: &mut Connection, credentials: Credentials) -> NirvResult<()>;
128    
129    /// Handle a query from the client and return a response
130    async fn handle_query(&self, conn: &Connection, query: ProtocolQuery) -> NirvResult<ProtocolResponse>;
131    
132    /// Get the protocol type this adapter handles
133    fn get_protocol_type(&self) -> ProtocolType;
134    
135    /// Parse protocol-specific message into internal representation
136    async fn parse_message(&self, conn: &Connection, data: &[u8]) -> NirvResult<ProtocolQuery>;
137    
138    /// Format internal query result into protocol-specific response
139    async fn format_response(&self, conn: &Connection, result: QueryResult) -> NirvResult<Vec<u8>>;
140    
141    /// Handle connection termination
142    async fn terminate_connection(&self, conn: &mut Connection) -> NirvResult<()>;
143}