sql_cli/data/
datatable.rs

1use crate::api_client::QueryResponse;
2use crate::data::data_provider::DataProvider;
3use crate::data::type_inference::{InferredType, TypeInference};
4use serde::{Deserialize, Serialize};
5use serde_json::Value as JsonValue;
6use std::collections::HashMap;
7use std::fmt;
8use std::sync::Arc;
9use tracing::debug;
10
11/// Represents the data type of a column
12#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
13pub enum DataType {
14    String,
15    Integer,
16    Float,
17    Boolean,
18    DateTime,
19    Null,
20    Mixed, // For columns with mixed types
21}
22
23impl DataType {
24    /// Infer type from a string value
25    pub fn infer_from_string(value: &str) -> Self {
26        // Handle explicit null string
27        if value.eq_ignore_ascii_case("null") {
28            return DataType::Null;
29        }
30
31        // Use the shared type inference logic
32        match TypeInference::infer_from_string(value) {
33            InferredType::Null => DataType::Null,
34            InferredType::Boolean => DataType::Boolean,
35            InferredType::Integer => DataType::Integer,
36            InferredType::Float => DataType::Float,
37            InferredType::DateTime => DataType::DateTime,
38            InferredType::String => DataType::String,
39        }
40    }
41
42    /// Check if a string looks like a datetime value
43    /// Delegates to shared type inference logic
44    fn looks_like_datetime(value: &str) -> bool {
45        TypeInference::looks_like_datetime(value)
46    }
47
48    /// Merge two types (for columns with mixed types)
49    pub fn merge(&self, other: &DataType) -> DataType {
50        if self == other {
51            return self.clone();
52        }
53
54        match (self, other) {
55            (DataType::Null, t) | (t, DataType::Null) => t.clone(),
56            (DataType::Integer, DataType::Float) | (DataType::Float, DataType::Integer) => {
57                DataType::Float
58            }
59            _ => DataType::Mixed,
60        }
61    }
62}
63
64/// Column metadata and definition
65#[derive(Debug, Clone, Serialize, Deserialize)]
66pub struct DataColumn {
67    pub name: String,
68    pub data_type: DataType,
69    pub nullable: bool,
70    pub unique_values: Option<usize>,
71    pub null_count: usize,
72    pub metadata: HashMap<String, String>,
73}
74
75impl DataColumn {
76    pub fn new(name: impl Into<String>) -> Self {
77        Self {
78            name: name.into(),
79            data_type: DataType::String,
80            nullable: true,
81            unique_values: None,
82            null_count: 0,
83            metadata: HashMap::new(),
84        }
85    }
86
87    pub fn with_type(mut self, data_type: DataType) -> Self {
88        self.data_type = data_type;
89        self
90    }
91
92    pub fn with_nullable(mut self, nullable: bool) -> Self {
93        self.nullable = nullable;
94        self
95    }
96}
97
98/// A single cell value in the table
99#[derive(Debug, Clone, PartialEq)]
100pub enum DataValue {
101    String(String),
102    InternedString(Arc<String>), // For repeated strings (e.g., status, trader names)
103    Integer(i64),
104    Float(f64),
105    Boolean(bool),
106    DateTime(String), // Store as ISO 8601 string for now
107    Null,
108}
109
110impl DataValue {
111    pub fn from_string(s: &str, data_type: &DataType) -> Self {
112        if s.is_empty() || s.eq_ignore_ascii_case("null") {
113            return DataValue::Null;
114        }
115
116        match data_type {
117            DataType::String => DataValue::String(s.to_string()),
118            DataType::Integer => s
119                .parse::<i64>()
120                .map(DataValue::Integer)
121                .unwrap_or_else(|_| DataValue::String(s.to_string())),
122            DataType::Float => s
123                .parse::<f64>()
124                .map(DataValue::Float)
125                .unwrap_or_else(|_| DataValue::String(s.to_string())),
126            DataType::Boolean => {
127                let lower = s.to_lowercase();
128                DataValue::Boolean(lower == "true" || lower == "1" || lower == "yes")
129            }
130            DataType::DateTime => DataValue::DateTime(s.to_string()),
131            DataType::Null => DataValue::Null,
132            DataType::Mixed => {
133                // Try to infer for mixed columns
134                let inferred = DataType::infer_from_string(s);
135                Self::from_string(s, &inferred)
136            }
137        }
138    }
139
140    pub fn is_null(&self) -> bool {
141        matches!(self, DataValue::Null)
142    }
143
144    pub fn data_type(&self) -> DataType {
145        match self {
146            DataValue::String(_) | DataValue::InternedString(_) => DataType::String,
147            DataValue::Integer(_) => DataType::Integer,
148            DataValue::Float(_) => DataType::Float,
149            DataValue::Boolean(_) => DataType::Boolean,
150            DataValue::DateTime(_) => DataType::DateTime,
151            DataValue::Null => DataType::Null,
152        }
153    }
154
155    /// Get string representation without allocation when possible
156    /// Returns owned String for compatibility but tries to reuse existing strings
157    pub fn to_string_optimized(&self) -> String {
158        match self {
159            DataValue::String(s) => s.clone(), // Clone existing string
160            DataValue::InternedString(s) => s.as_ref().clone(), // Clone from Rc
161            DataValue::DateTime(s) => s.clone(), // Clone existing string
162            DataValue::Integer(i) => i.to_string(),
163            DataValue::Float(f) => f.to_string(),
164            DataValue::Boolean(b) => {
165                if *b {
166                    "true".to_string()
167                } else {
168                    "false".to_string()
169                }
170            }
171            DataValue::Null => String::new(), // Empty string, minimal allocation
172        }
173    }
174}
175
176impl fmt::Display for DataValue {
177    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
178        match self {
179            DataValue::String(s) => write!(f, "{}", s),
180            DataValue::InternedString(s) => write!(f, "{}", s),
181            DataValue::Integer(i) => write!(f, "{}", i),
182            DataValue::Float(fl) => write!(f, "{}", fl),
183            DataValue::Boolean(b) => write!(f, "{}", b),
184            DataValue::DateTime(dt) => write!(f, "{}", dt),
185            DataValue::Null => write!(f, ""),
186        }
187    }
188}
189
190/// A row of data in the table
191#[derive(Debug, Clone)]
192pub struct DataRow {
193    pub values: Vec<DataValue>,
194}
195
196impl DataRow {
197    pub fn new(values: Vec<DataValue>) -> Self {
198        Self { values }
199    }
200
201    pub fn get(&self, index: usize) -> Option<&DataValue> {
202        self.values.get(index)
203    }
204
205    pub fn get_mut(&mut self, index: usize) -> Option<&mut DataValue> {
206        self.values.get_mut(index)
207    }
208
209    pub fn len(&self) -> usize {
210        self.values.len()
211    }
212
213    pub fn is_empty(&self) -> bool {
214        self.values.is_empty()
215    }
216}
217
218/// The main DataTable structure
219#[derive(Debug, Clone)]
220pub struct DataTable {
221    pub name: String,
222    pub columns: Vec<DataColumn>,
223    pub rows: Vec<DataRow>,
224    pub metadata: HashMap<String, String>,
225}
226
227impl DataTable {
228    pub fn new(name: impl Into<String>) -> Self {
229        Self {
230            name: name.into(),
231            columns: Vec::new(),
232            rows: Vec::new(),
233            metadata: HashMap::new(),
234        }
235    }
236
237    pub fn add_column(&mut self, column: DataColumn) -> &mut Self {
238        self.columns.push(column);
239        self
240    }
241
242    pub fn add_row(&mut self, row: DataRow) -> Result<(), String> {
243        if row.len() != self.columns.len() {
244            return Err(format!(
245                "Row has {} values but table has {} columns",
246                row.len(),
247                self.columns.len()
248            ));
249        }
250        self.rows.push(row);
251        Ok(())
252    }
253
254    pub fn get_column(&self, name: &str) -> Option<&DataColumn> {
255        self.columns.iter().find(|c| c.name == name)
256    }
257
258    pub fn get_column_index(&self, name: &str) -> Option<usize> {
259        self.columns.iter().position(|c| c.name == name)
260    }
261
262    pub fn column_count(&self) -> usize {
263        self.columns.len()
264    }
265
266    pub fn row_count(&self) -> usize {
267        self.rows.len()
268    }
269
270    pub fn is_empty(&self) -> bool {
271        self.rows.is_empty()
272    }
273
274    /// Get column names as a vector
275    pub fn column_names(&self) -> Vec<String> {
276        self.columns.iter().map(|c| c.name.clone()).collect()
277    }
278
279    /// Infer and update column types based on data
280    pub fn infer_column_types(&mut self) {
281        for (col_idx, column) in self.columns.iter_mut().enumerate() {
282            let mut inferred_type = DataType::Null;
283            let mut null_count = 0;
284            let mut unique_values = std::collections::HashSet::new();
285
286            for row in &self.rows {
287                if let Some(value) = row.get(col_idx) {
288                    if value.is_null() {
289                        null_count += 1;
290                    } else {
291                        let value_type = value.data_type();
292                        inferred_type = inferred_type.merge(&value_type);
293                        unique_values.insert(value.to_string());
294                    }
295                }
296            }
297
298            column.data_type = inferred_type;
299            column.null_count = null_count;
300            column.nullable = null_count > 0;
301            column.unique_values = Some(unique_values.len());
302        }
303    }
304
305    /// Get a value at specific row and column
306    pub fn get_value(&self, row: usize, col: usize) -> Option<&DataValue> {
307        self.rows.get(row)?.get(col)
308    }
309
310    /// Get a value by row index and column name
311    pub fn get_value_by_name(&self, row: usize, col_name: &str) -> Option<&DataValue> {
312        let col_idx = self.get_column_index(col_name)?;
313        self.get_value(row, col_idx)
314    }
315
316    /// Convert to a vector of string vectors (for display/compatibility)
317    pub fn to_string_table(&self) -> Vec<Vec<String>> {
318        self.rows
319            .iter()
320            .map(|row| row.values.iter().map(|v| v.to_string_optimized()).collect())
321            .collect()
322    }
323
324    /// Get table statistics
325    pub fn get_stats(&self) -> DataTableStats {
326        DataTableStats {
327            row_count: self.row_count(),
328            column_count: self.column_count(),
329            memory_size: self.estimate_memory_size(),
330            null_count: self.columns.iter().map(|c| c.null_count).sum(),
331        }
332    }
333
334    /// Generate a debug dump string for display
335    pub fn debug_dump(&self) -> String {
336        let mut output = String::new();
337
338        output.push_str(&format!("DataTable: {}\n", self.name));
339        output.push_str(&format!(
340            "Rows: {} | Columns: {}\n",
341            self.row_count(),
342            self.column_count()
343        ));
344
345        if !self.metadata.is_empty() {
346            output.push_str("Metadata:\n");
347            for (key, value) in &self.metadata {
348                output.push_str(&format!("  {}: {}\n", key, value));
349            }
350        }
351
352        output.push_str("\nColumns:\n");
353        for column in &self.columns {
354            output.push_str(&format!("  {} ({:?})", column.name, column.data_type));
355            if column.nullable {
356                output.push_str(&format!(" - nullable, {} nulls", column.null_count));
357            }
358            if let Some(unique) = column.unique_values {
359                output.push_str(&format!(", {} unique", unique));
360            }
361            output.push('\n');
362        }
363
364        // Show first few rows
365        if self.row_count() > 0 {
366            let sample_size = 5.min(self.row_count());
367            output.push_str(&format!("\nFirst {} rows:\n", sample_size));
368
369            for row_idx in 0..sample_size {
370                output.push_str(&format!("  [{}]: ", row_idx));
371                for (col_idx, value) in self.rows[row_idx].values.iter().enumerate() {
372                    if col_idx > 0 {
373                        output.push_str(", ");
374                    }
375                    output.push_str(&value.to_string());
376                }
377                output.push('\n');
378            }
379        }
380
381        output
382    }
383
384    pub fn estimate_memory_size(&self) -> usize {
385        // Base structure size
386        let mut size = std::mem::size_of::<Self>();
387
388        // Column metadata
389        size += self.columns.len() * std::mem::size_of::<DataColumn>();
390        for col in &self.columns {
391            size += col.name.len();
392        }
393
394        // Row structure overhead
395        size += self.rows.len() * std::mem::size_of::<DataRow>();
396
397        // Actual data values
398        for row in &self.rows {
399            for value in &row.values {
400                // Base enum size
401                size += std::mem::size_of::<DataValue>();
402                // Add string content size
403                match value {
404                    DataValue::String(s) | DataValue::DateTime(s) => size += s.len(),
405                    _ => {} // Numbers and booleans are inline
406                }
407            }
408        }
409
410        size
411    }
412
413    /// V46: Create DataTable from QueryResponse
414    /// This is the key conversion function that bridges old and new systems
415    pub fn from_query_response(response: &QueryResponse, table_name: &str) -> Result<Self, String> {
416        debug!(
417            "V46: Converting QueryResponse to DataTable for table '{}'",
418            table_name
419        );
420
421        // Track memory before conversion
422        crate::utils::memory_tracker::track_memory("start_from_query_response");
423
424        let mut table = DataTable::new(table_name);
425
426        // Extract column names and types from first row
427        if let Some(first_row) = response.data.first() {
428            if let Some(obj) = first_row.as_object() {
429                // Create columns based on the keys in the JSON object
430                for key in obj.keys() {
431                    let column = DataColumn::new(key.clone());
432                    table.add_column(column);
433                }
434
435                // Now convert all rows
436                for json_row in &response.data {
437                    if let Some(row_obj) = json_row.as_object() {
438                        let mut values = Vec::new();
439
440                        // Ensure we get values in the same order as columns
441                        for column in &table.columns {
442                            let value = row_obj
443                                .get(&column.name)
444                                .map(|v| json_value_to_data_value(v))
445                                .unwrap_or(DataValue::Null);
446                            values.push(value);
447                        }
448
449                        table.add_row(DataRow::new(values))?;
450                    }
451                }
452
453                // Infer column types from the data
454                table.infer_column_types();
455
456                // Add metadata
457                if let Some(source) = &response.source {
458                    table.metadata.insert("source".to_string(), source.clone());
459                }
460                if let Some(cached) = response.cached {
461                    table
462                        .metadata
463                        .insert("cached".to_string(), cached.to_string());
464                }
465                table
466                    .metadata
467                    .insert("original_count".to_string(), response.count.to_string());
468
469                debug!(
470                    "V46: Created DataTable with {} columns and {} rows",
471                    table.column_count(),
472                    table.row_count()
473                );
474            } else {
475                // Handle non-object JSON (single values)
476                table.add_column(DataColumn::new("value"));
477                for json_value in &response.data {
478                    let value = json_value_to_data_value(json_value);
479                    table.add_row(DataRow::new(vec![value]))?;
480                }
481            }
482        }
483
484        Ok(table)
485    }
486
487    /// Get a single row by index
488    pub fn get_row(&self, index: usize) -> Option<&DataRow> {
489        self.rows.get(index)
490    }
491
492    /// V50: Get a single row as strings
493    pub fn get_row_as_strings(&self, index: usize) -> Option<Vec<String>> {
494        self.rows.get(index).map(|row| {
495            row.values
496                .iter()
497                .map(|value| value.to_string_optimized())
498                .collect()
499        })
500    }
501
502    /// Pretty print the DataTable with a nice box drawing
503    pub fn pretty_print(&self) -> String {
504        let mut output = String::new();
505
506        // Header
507        output.push_str("╔═══════════════════════════════════════════════════════╗\n");
508        output.push_str(&format!("║ DataTable: {:^41} ║\n", self.name));
509        output.push_str("╠═══════════════════════════════════════════════════════╣\n");
510
511        // Summary stats
512        output.push_str(&format!(
513            "║ Rows: {:6} | Columns: {:3} | Memory: ~{:6} bytes ║\n",
514            self.row_count(),
515            self.column_count(),
516            self.get_stats().memory_size
517        ));
518
519        // Metadata if any
520        if !self.metadata.is_empty() {
521            output.push_str("╠═══════════════════════════════════════════════════════╣\n");
522            output.push_str("║ Metadata:                                             ║\n");
523            for (key, value) in &self.metadata {
524                let truncated_value = if value.len() > 35 {
525                    format!("{}...", &value[..32])
526                } else {
527                    value.clone()
528                };
529                output.push_str(&format!(
530                    "║   {:15} : {:35} ║\n",
531                    Self::truncate_string(key, 15),
532                    truncated_value
533                ));
534            }
535        }
536
537        // Column details
538        output.push_str("╠═══════════════════════════════════════════════════════╣\n");
539        output.push_str("║ Columns:                                              ║\n");
540        output.push_str("╟───────────────────┬──────────┬─────────┬──────┬──────╢\n");
541        output.push_str("║ Name              │ Type     │ Nullable│ Nulls│Unique║\n");
542        output.push_str("╟───────────────────┼──────────┼─────────┼──────┼──────╢\n");
543
544        for column in &self.columns {
545            let type_str = match &column.data_type {
546                DataType::String => "String",
547                DataType::Integer => "Integer",
548                DataType::Float => "Float",
549                DataType::Boolean => "Boolean",
550                DataType::DateTime => "DateTime",
551                DataType::Null => "Null",
552                DataType::Mixed => "Mixed",
553            };
554
555            output.push_str(&format!(
556                "║ {:17} │ {:8} │ {:7} │ {:4} │ {:4} ║\n",
557                Self::truncate_string(&column.name, 17),
558                type_str,
559                if column.nullable { "Yes" } else { "No" },
560                column.null_count,
561                column.unique_values.unwrap_or(0)
562            ));
563        }
564
565        output.push_str("╚═══════════════════════════════════════════════════════╝\n");
566
567        // Sample data (first 5 rows)
568        output.push_str("\nSample Data (first 5 rows):\n");
569        let sample_count = self.rows.len().min(5);
570
571        if sample_count > 0 {
572            // Column headers
573            output.push_str("┌");
574            for (i, col) in self.columns.iter().enumerate() {
575                if i > 0 {
576                    output.push_str("┬");
577                }
578                output.push_str(&"─".repeat(20));
579            }
580            output.push_str("┐\n");
581
582            output.push_str("│");
583            for col in &self.columns {
584                output.push_str(&format!(" {:^18} │", Self::truncate_string(&col.name, 18)));
585            }
586            output.push_str("\n");
587
588            output.push_str("├");
589            for (i, _) in self.columns.iter().enumerate() {
590                if i > 0 {
591                    output.push_str("┼");
592                }
593                output.push_str(&"─".repeat(20));
594            }
595            output.push_str("┤\n");
596
597            // Data rows
598            for row_idx in 0..sample_count {
599                if let Some(row) = self.rows.get(row_idx) {
600                    output.push_str("│");
601                    for value in &row.values {
602                        let value_str = value.to_string();
603                        output
604                            .push_str(&format!(" {:18} │", Self::truncate_string(&value_str, 18)));
605                    }
606                    output.push_str("\n");
607                }
608            }
609
610            output.push_str("└");
611            for (i, _) in self.columns.iter().enumerate() {
612                if i > 0 {
613                    output.push_str("┴");
614                }
615                output.push_str(&"─".repeat(20));
616            }
617            output.push_str("┘\n");
618        }
619
620        output
621    }
622
623    fn truncate_string(s: &str, max_len: usize) -> String {
624        if s.len() > max_len {
625            format!("{}...", &s[..max_len - 3])
626        } else {
627            s.to_string()
628        }
629    }
630
631    /// Get a schema summary of the DataTable
632    pub fn get_schema_summary(&self) -> String {
633        let mut summary = String::new();
634        summary.push_str(&format!(
635            "DataTable Schema ({} columns, {} rows):\n",
636            self.columns.len(),
637            self.rows.len()
638        ));
639
640        for (idx, column) in self.columns.iter().enumerate() {
641            let type_str = match &column.data_type {
642                DataType::String => "String",
643                DataType::Integer => "Integer",
644                DataType::Float => "Float",
645                DataType::Boolean => "Boolean",
646                DataType::DateTime => "DateTime",
647                DataType::Null => "Null",
648                DataType::Mixed => "Mixed",
649            };
650
651            let nullable_str = if column.nullable {
652                "nullable"
653            } else {
654                "not null"
655            };
656            let null_info = if column.null_count > 0 {
657                format!(", {} nulls", column.null_count)
658            } else {
659                String::new()
660            };
661
662            summary.push_str(&format!(
663                "  [{:3}] {} : {} ({}{})\n",
664                idx, column.name, type_str, nullable_str, null_info
665            ));
666        }
667
668        summary
669    }
670
671    /// Get detailed schema information as a structured format
672    pub fn get_schema_info(&self) -> Vec<(String, String, bool, usize)> {
673        self.columns
674            .iter()
675            .map(|col| {
676                let type_name = format!("{:?}", col.data_type);
677                (col.name.clone(), type_name, col.nullable, col.null_count)
678            })
679            .collect()
680    }
681
682    /// Reserve capacity for rows to avoid reallocations
683    pub fn reserve_rows(&mut self, additional: usize) {
684        self.rows.reserve(additional);
685    }
686
687    /// Shrink vectors to fit actual data (removes excess capacity)
688    pub fn shrink_to_fit(&mut self) {
689        self.rows.shrink_to_fit();
690        for column in &mut self.columns {
691            // Shrink any column-specific data if needed
692        }
693    }
694
695    /// Get actual memory usage estimate (more accurate than estimate_memory_size)
696    pub fn get_memory_usage(&self) -> usize {
697        let mut size = std::mem::size_of::<Self>();
698
699        // Account for string allocations
700        size += self.name.capacity();
701
702        // Account for columns
703        size += self.columns.capacity() * std::mem::size_of::<DataColumn>();
704        for col in &self.columns {
705            size += col.name.capacity();
706        }
707
708        // Account for rows and their capacity
709        size += self.rows.capacity() * std::mem::size_of::<DataRow>();
710
711        // Account for actual data values
712        for row in &self.rows {
713            size += row.values.capacity() * std::mem::size_of::<DataValue>();
714            for value in &row.values {
715                match value {
716                    DataValue::String(s) => size += s.capacity(),
717                    DataValue::InternedString(_) => size += std::mem::size_of::<Arc<String>>(),
718                    DataValue::DateTime(s) => size += s.capacity(),
719                    _ => {} // Other types are inline
720                }
721            }
722        }
723
724        // Account for metadata
725        size += self.metadata.capacity() * std::mem::size_of::<(String, String)>();
726        for (k, v) in &self.metadata {
727            size += k.capacity() + v.capacity();
728        }
729
730        size
731    }
732}
733
734/// V46: Helper function to convert JSON value to DataValue
735fn json_value_to_data_value(json: &JsonValue) -> DataValue {
736    match json {
737        JsonValue::Null => DataValue::Null,
738        JsonValue::Bool(b) => DataValue::Boolean(*b),
739        JsonValue::Number(n) => {
740            if let Some(i) = n.as_i64() {
741                DataValue::Integer(i)
742            } else if let Some(f) = n.as_f64() {
743                DataValue::Float(f)
744            } else {
745                DataValue::String(n.to_string())
746            }
747        }
748        JsonValue::String(s) => {
749            // Try to detect if it's a date/time
750            if s.contains('-') && s.len() >= 8 && s.len() <= 30 {
751                // Simple heuristic for dates
752                DataValue::DateTime(s.clone())
753            } else {
754                DataValue::String(s.clone())
755            }
756        }
757        JsonValue::Array(_) | JsonValue::Object(_) => {
758            // Store complex types as JSON string
759            DataValue::String(json.to_string())
760        }
761    }
762}
763
764/// Statistics about a DataTable
765#[derive(Debug, Clone)]
766pub struct DataTableStats {
767    pub row_count: usize,
768    pub column_count: usize,
769    pub memory_size: usize,
770    pub null_count: usize,
771}
772
773/// Implementation of DataProvider for DataTable
774/// This allows DataTable to be used wherever DataProvider trait is expected
775impl DataProvider for DataTable {
776    fn get_row(&self, index: usize) -> Option<Vec<String>> {
777        self.rows
778            .get(index)
779            .map(|row| row.values.iter().map(|v| v.to_string_optimized()).collect())
780    }
781
782    fn get_column_names(&self) -> Vec<String> {
783        self.column_names()
784    }
785
786    fn get_row_count(&self) -> usize {
787        self.row_count()
788    }
789
790    fn get_column_count(&self) -> usize {
791        self.column_count()
792    }
793}
794
795#[cfg(test)]
796mod tests {
797    use super::*;
798
799    #[test]
800    fn test_data_type_inference() {
801        assert_eq!(DataType::infer_from_string("123"), DataType::Integer);
802        assert_eq!(DataType::infer_from_string("123.45"), DataType::Float);
803        assert_eq!(DataType::infer_from_string("true"), DataType::Boolean);
804        assert_eq!(DataType::infer_from_string("hello"), DataType::String);
805        assert_eq!(DataType::infer_from_string(""), DataType::Null);
806        assert_eq!(
807            DataType::infer_from_string("2024-01-01"),
808            DataType::DateTime
809        );
810    }
811
812    #[test]
813    fn test_datatable_creation() {
814        let mut table = DataTable::new("test");
815
816        table.add_column(DataColumn::new("id").with_type(DataType::Integer));
817        table.add_column(DataColumn::new("name").with_type(DataType::String));
818        table.add_column(DataColumn::new("active").with_type(DataType::Boolean));
819
820        assert_eq!(table.column_count(), 3);
821        assert_eq!(table.row_count(), 0);
822
823        let row = DataRow::new(vec![
824            DataValue::Integer(1),
825            DataValue::String("Alice".to_string()),
826            DataValue::Boolean(true),
827        ]);
828
829        table.add_row(row).unwrap();
830        assert_eq!(table.row_count(), 1);
831
832        let value = table.get_value_by_name(0, "name").unwrap();
833        assert_eq!(value.to_string(), "Alice");
834    }
835
836    #[test]
837    fn test_type_inference() {
838        let mut table = DataTable::new("test");
839
840        // Add columns without types
841        table.add_column(DataColumn::new("mixed"));
842
843        // Add rows with different types
844        table
845            .add_row(DataRow::new(vec![DataValue::Integer(1)]))
846            .unwrap();
847        table
848            .add_row(DataRow::new(vec![DataValue::Float(2.5)]))
849            .unwrap();
850        table.add_row(DataRow::new(vec![DataValue::Null])).unwrap();
851
852        table.infer_column_types();
853
854        // Should infer Float since we have both Integer and Float
855        assert_eq!(table.columns[0].data_type, DataType::Float);
856        assert_eq!(table.columns[0].null_count, 1);
857        assert!(table.columns[0].nullable);
858    }
859
860    #[test]
861    fn test_from_query_response() {
862        use crate::api_client::{QueryInfo, QueryResponse};
863        use serde_json::json;
864
865        let response = QueryResponse {
866            query: QueryInfo {
867                select: vec!["id".to_string(), "name".to_string(), "age".to_string()],
868                where_clause: None,
869                order_by: None,
870            },
871            data: vec![
872                json!({
873                    "id": 1,
874                    "name": "Alice",
875                    "age": 30
876                }),
877                json!({
878                    "id": 2,
879                    "name": "Bob",
880                    "age": 25
881                }),
882                json!({
883                    "id": 3,
884                    "name": "Carol",
885                    "age": null
886                }),
887            ],
888            count: 3,
889            source: Some("test.csv".to_string()),
890            table: Some("test".to_string()),
891            cached: Some(false),
892        };
893
894        let table = DataTable::from_query_response(&response, "test").unwrap();
895
896        assert_eq!(table.name, "test");
897        assert_eq!(table.row_count(), 3);
898        assert_eq!(table.column_count(), 3);
899
900        // Check column names
901        let col_names = table.column_names();
902        assert!(col_names.contains(&"id".to_string()));
903        assert!(col_names.contains(&"name".to_string()));
904        assert!(col_names.contains(&"age".to_string()));
905
906        // Check metadata
907        assert_eq!(table.metadata.get("source"), Some(&"test.csv".to_string()));
908        assert_eq!(table.metadata.get("cached"), Some(&"false".to_string()));
909
910        // Check first row values
911        assert_eq!(
912            table.get_value_by_name(0, "id"),
913            Some(&DataValue::Integer(1))
914        );
915        assert_eq!(
916            table.get_value_by_name(0, "name"),
917            Some(&DataValue::String("Alice".to_string()))
918        );
919        assert_eq!(
920            table.get_value_by_name(0, "age"),
921            Some(&DataValue::Integer(30))
922        );
923
924        // Check null handling
925        assert_eq!(table.get_value_by_name(2, "age"), Some(&DataValue::Null));
926    }
927}
sql_cli/data/datatable.rs

sql_cli/data/
datatable.rs
