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, PartialOrd)]
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    /// Create a DUAL table (similar to Oracle's DUAL) with one row and one column
238    /// Used for evaluating expressions without a data source
239    pub fn dual() -> Self {
240        let mut table = DataTable::new("DUAL");
241        table.add_column(DataColumn::new("DUMMY").with_type(DataType::String));
242        table
243            .add_row(DataRow::new(vec![DataValue::String("X".to_string())]))
244            .unwrap();
245        table
246    }
247
248    pub fn add_column(&mut self, column: DataColumn) -> &mut Self {
249        self.columns.push(column);
250        self
251    }
252
253    pub fn add_row(&mut self, row: DataRow) -> Result<(), String> {
254        if row.len() != self.columns.len() {
255            return Err(format!(
256                "Row has {} values but table has {} columns",
257                row.len(),
258                self.columns.len()
259            ));
260        }
261        self.rows.push(row);
262        Ok(())
263    }
264
265    pub fn get_column(&self, name: &str) -> Option<&DataColumn> {
266        self.columns.iter().find(|c| c.name == name)
267    }
268
269    pub fn get_column_index(&self, name: &str) -> Option<usize> {
270        self.columns.iter().position(|c| c.name == name)
271    }
272
273    pub fn column_count(&self) -> usize {
274        self.columns.len()
275    }
276
277    pub fn row_count(&self) -> usize {
278        self.rows.len()
279    }
280
281    pub fn is_empty(&self) -> bool {
282        self.rows.is_empty()
283    }
284
285    /// Get column names as a vector
286    pub fn column_names(&self) -> Vec<String> {
287        self.columns.iter().map(|c| c.name.clone()).collect()
288    }
289
290    /// Infer and update column types based on data
291    pub fn infer_column_types(&mut self) {
292        for (col_idx, column) in self.columns.iter_mut().enumerate() {
293            let mut inferred_type = DataType::Null;
294            let mut null_count = 0;
295            let mut unique_values = std::collections::HashSet::new();
296
297            for row in &self.rows {
298                if let Some(value) = row.get(col_idx) {
299                    if value.is_null() {
300                        null_count += 1;
301                    } else {
302                        let value_type = value.data_type();
303                        inferred_type = inferred_type.merge(&value_type);
304                        unique_values.insert(value.to_string());
305                    }
306                }
307            }
308
309            column.data_type = inferred_type;
310            column.null_count = null_count;
311            column.nullable = null_count > 0;
312            column.unique_values = Some(unique_values.len());
313        }
314    }
315
316    /// Get a value at specific row and column
317    pub fn get_value(&self, row: usize, col: usize) -> Option<&DataValue> {
318        self.rows.get(row)?.get(col)
319    }
320
321    /// Get a value by row index and column name
322    pub fn get_value_by_name(&self, row: usize, col_name: &str) -> Option<&DataValue> {
323        let col_idx = self.get_column_index(col_name)?;
324        self.get_value(row, col_idx)
325    }
326
327    /// Convert to a vector of string vectors (for display/compatibility)
328    pub fn to_string_table(&self) -> Vec<Vec<String>> {
329        self.rows
330            .iter()
331            .map(|row| row.values.iter().map(|v| v.to_string_optimized()).collect())
332            .collect()
333    }
334
335    /// Get table statistics
336    pub fn get_stats(&self) -> DataTableStats {
337        DataTableStats {
338            row_count: self.row_count(),
339            column_count: self.column_count(),
340            memory_size: self.estimate_memory_size(),
341            null_count: self.columns.iter().map(|c| c.null_count).sum(),
342        }
343    }
344
345    /// Generate a debug dump string for display
346    pub fn debug_dump(&self) -> String {
347        let mut output = String::new();
348
349        output.push_str(&format!("DataTable: {}\n", self.name));
350        output.push_str(&format!(
351            "Rows: {} | Columns: {}\n",
352            self.row_count(),
353            self.column_count()
354        ));
355
356        if !self.metadata.is_empty() {
357            output.push_str("Metadata:\n");
358            for (key, value) in &self.metadata {
359                output.push_str(&format!("  {}: {}\n", key, value));
360            }
361        }
362
363        output.push_str("\nColumns:\n");
364        for column in &self.columns {
365            output.push_str(&format!("  {} ({:?})", column.name, column.data_type));
366            if column.nullable {
367                output.push_str(&format!(" - nullable, {} nulls", column.null_count));
368            }
369            if let Some(unique) = column.unique_values {
370                output.push_str(&format!(", {} unique", unique));
371            }
372            output.push('\n');
373        }
374
375        // Show first few rows
376        if self.row_count() > 0 {
377            let sample_size = 5.min(self.row_count());
378            output.push_str(&format!("\nFirst {} rows:\n", sample_size));
379
380            for row_idx in 0..sample_size {
381                output.push_str(&format!("  [{}]: ", row_idx));
382                for (col_idx, value) in self.rows[row_idx].values.iter().enumerate() {
383                    if col_idx > 0 {
384                        output.push_str(", ");
385                    }
386                    output.push_str(&value.to_string());
387                }
388                output.push('\n');
389            }
390        }
391
392        output
393    }
394
395    pub fn estimate_memory_size(&self) -> usize {
396        // Base structure size
397        let mut size = std::mem::size_of::<Self>();
398
399        // Column metadata
400        size += self.columns.len() * std::mem::size_of::<DataColumn>();
401        for col in &self.columns {
402            size += col.name.len();
403        }
404
405        // Row structure overhead
406        size += self.rows.len() * std::mem::size_of::<DataRow>();
407
408        // Actual data values
409        for row in &self.rows {
410            for value in &row.values {
411                // Base enum size
412                size += std::mem::size_of::<DataValue>();
413                // Add string content size
414                match value {
415                    DataValue::String(s) | DataValue::DateTime(s) => size += s.len(),
416                    _ => {} // Numbers and booleans are inline
417                }
418            }
419        }
420
421        size
422    }
423
424    /// V46: Create DataTable from QueryResponse
425    /// This is the key conversion function that bridges old and new systems
426    pub fn from_query_response(response: &QueryResponse, table_name: &str) -> Result<Self, String> {
427        debug!(
428            "V46: Converting QueryResponse to DataTable for table '{}'",
429            table_name
430        );
431
432        // Track memory before conversion
433        crate::utils::memory_tracker::track_memory("start_from_query_response");
434
435        let mut table = DataTable::new(table_name);
436
437        // Extract column names and types from first row
438        if let Some(first_row) = response.data.first() {
439            if let Some(obj) = first_row.as_object() {
440                // Create columns based on the keys in the JSON object
441                for key in obj.keys() {
442                    let column = DataColumn::new(key.clone());
443                    table.add_column(column);
444                }
445
446                // Now convert all rows
447                for json_row in &response.data {
448                    if let Some(row_obj) = json_row.as_object() {
449                        let mut values = Vec::new();
450
451                        // Ensure we get values in the same order as columns
452                        for column in &table.columns {
453                            let value = row_obj
454                                .get(&column.name)
455                                .map(|v| json_value_to_data_value(v))
456                                .unwrap_or(DataValue::Null);
457                            values.push(value);
458                        }
459
460                        table.add_row(DataRow::new(values))?;
461                    }
462                }
463
464                // Infer column types from the data
465                table.infer_column_types();
466
467                // Add metadata
468                if let Some(source) = &response.source {
469                    table.metadata.insert("source".to_string(), source.clone());
470                }
471                if let Some(cached) = response.cached {
472                    table
473                        .metadata
474                        .insert("cached".to_string(), cached.to_string());
475                }
476                table
477                    .metadata
478                    .insert("original_count".to_string(), response.count.to_string());
479
480                debug!(
481                    "V46: Created DataTable with {} columns and {} rows",
482                    table.column_count(),
483                    table.row_count()
484                );
485            } else {
486                // Handle non-object JSON (single values)
487                table.add_column(DataColumn::new("value"));
488                for json_value in &response.data {
489                    let value = json_value_to_data_value(json_value);
490                    table.add_row(DataRow::new(vec![value]))?;
491                }
492            }
493        }
494
495        Ok(table)
496    }
497
498    /// Get a single row by index
499    pub fn get_row(&self, index: usize) -> Option<&DataRow> {
500        self.rows.get(index)
501    }
502
503    /// V50: Get a single row as strings
504    pub fn get_row_as_strings(&self, index: usize) -> Option<Vec<String>> {
505        self.rows.get(index).map(|row| {
506            row.values
507                .iter()
508                .map(|value| value.to_string_optimized())
509                .collect()
510        })
511    }
512
513    /// Pretty print the DataTable with a nice box drawing
514    pub fn pretty_print(&self) -> String {
515        let mut output = String::new();
516
517        // Header
518        output.push_str("╔═══════════════════════════════════════════════════════╗\n");
519        output.push_str(&format!("║ DataTable: {:^41} ║\n", self.name));
520        output.push_str("╠═══════════════════════════════════════════════════════╣\n");
521
522        // Summary stats
523        output.push_str(&format!(
524            "║ Rows: {:6} | Columns: {:3} | Memory: ~{:6} bytes ║\n",
525            self.row_count(),
526            self.column_count(),
527            self.get_stats().memory_size
528        ));
529
530        // Metadata if any
531        if !self.metadata.is_empty() {
532            output.push_str("╠═══════════════════════════════════════════════════════╣\n");
533            output.push_str("║ Metadata:                                             ║\n");
534            for (key, value) in &self.metadata {
535                let truncated_value = if value.len() > 35 {
536                    format!("{}...", &value[..32])
537                } else {
538                    value.clone()
539                };
540                output.push_str(&format!(
541                    "║   {:15} : {:35} ║\n",
542                    Self::truncate_string(key, 15),
543                    truncated_value
544                ));
545            }
546        }
547
548        // Column details
549        output.push_str("╠═══════════════════════════════════════════════════════╣\n");
550        output.push_str("║ Columns:                                              ║\n");
551        output.push_str("╟───────────────────┬──────────┬─────────┬──────┬──────╢\n");
552        output.push_str("║ Name              │ Type     │ Nullable│ Nulls│Unique║\n");
553        output.push_str("╟───────────────────┼──────────┼─────────┼──────┼──────╢\n");
554
555        for column in &self.columns {
556            let type_str = match &column.data_type {
557                DataType::String => "String",
558                DataType::Integer => "Integer",
559                DataType::Float => "Float",
560                DataType::Boolean => "Boolean",
561                DataType::DateTime => "DateTime",
562                DataType::Null => "Null",
563                DataType::Mixed => "Mixed",
564            };
565
566            output.push_str(&format!(
567                "║ {:17} │ {:8} │ {:7} │ {:4} │ {:4} ║\n",
568                Self::truncate_string(&column.name, 17),
569                type_str,
570                if column.nullable { "Yes" } else { "No" },
571                column.null_count,
572                column.unique_values.unwrap_or(0)
573            ));
574        }
575
576        output.push_str("╚═══════════════════════════════════════════════════════╝\n");
577
578        // Sample data (first 5 rows)
579        output.push_str("\nSample Data (first 5 rows):\n");
580        let sample_count = self.rows.len().min(5);
581
582        if sample_count > 0 {
583            // Column headers
584            output.push_str("┌");
585            for (i, col) in self.columns.iter().enumerate() {
586                if i > 0 {
587                    output.push_str("┬");
588                }
589                output.push_str(&"─".repeat(20));
590            }
591            output.push_str("┐\n");
592
593            output.push_str("│");
594            for col in &self.columns {
595                output.push_str(&format!(" {:^18} │", Self::truncate_string(&col.name, 18)));
596            }
597            output.push_str("\n");
598
599            output.push_str("├");
600            for (i, _) in self.columns.iter().enumerate() {
601                if i > 0 {
602                    output.push_str("┼");
603                }
604                output.push_str(&"─".repeat(20));
605            }
606            output.push_str("┤\n");
607
608            // Data rows
609            for row_idx in 0..sample_count {
610                if let Some(row) = self.rows.get(row_idx) {
611                    output.push_str("│");
612                    for value in &row.values {
613                        let value_str = value.to_string();
614                        output
615                            .push_str(&format!(" {:18} │", Self::truncate_string(&value_str, 18)));
616                    }
617                    output.push_str("\n");
618                }
619            }
620
621            output.push_str("└");
622            for (i, _) in self.columns.iter().enumerate() {
623                if i > 0 {
624                    output.push_str("┴");
625                }
626                output.push_str(&"─".repeat(20));
627            }
628            output.push_str("┘\n");
629        }
630
631        output
632    }
633
634    fn truncate_string(s: &str, max_len: usize) -> String {
635        if s.len() > max_len {
636            format!("{}...", &s[..max_len - 3])
637        } else {
638            s.to_string()
639        }
640    }
641
642    /// Get a schema summary of the DataTable
643    pub fn get_schema_summary(&self) -> String {
644        let mut summary = String::new();
645        summary.push_str(&format!(
646            "DataTable Schema ({} columns, {} rows):\n",
647            self.columns.len(),
648            self.rows.len()
649        ));
650
651        for (idx, column) in self.columns.iter().enumerate() {
652            let type_str = match &column.data_type {
653                DataType::String => "String",
654                DataType::Integer => "Integer",
655                DataType::Float => "Float",
656                DataType::Boolean => "Boolean",
657                DataType::DateTime => "DateTime",
658                DataType::Null => "Null",
659                DataType::Mixed => "Mixed",
660            };
661
662            let nullable_str = if column.nullable {
663                "nullable"
664            } else {
665                "not null"
666            };
667            let null_info = if column.null_count > 0 {
668                format!(", {} nulls", column.null_count)
669            } else {
670                String::new()
671            };
672
673            summary.push_str(&format!(
674                "  [{:3}] {} : {} ({}{})\n",
675                idx, column.name, type_str, nullable_str, null_info
676            ));
677        }
678
679        summary
680    }
681
682    /// Get detailed schema information as a structured format
683    pub fn get_schema_info(&self) -> Vec<(String, String, bool, usize)> {
684        self.columns
685            .iter()
686            .map(|col| {
687                let type_name = format!("{:?}", col.data_type);
688                (col.name.clone(), type_name, col.nullable, col.null_count)
689            })
690            .collect()
691    }
692
693    /// Reserve capacity for rows to avoid reallocations
694    pub fn reserve_rows(&mut self, additional: usize) {
695        self.rows.reserve(additional);
696    }
697
698    /// Shrink vectors to fit actual data (removes excess capacity)
699    pub fn shrink_to_fit(&mut self) {
700        self.rows.shrink_to_fit();
701        for column in &mut self.columns {
702            // Shrink any column-specific data if needed
703        }
704    }
705
706    /// Get actual memory usage estimate (more accurate than estimate_memory_size)
707    pub fn get_memory_usage(&self) -> usize {
708        let mut size = std::mem::size_of::<Self>();
709
710        // Account for string allocations
711        size += self.name.capacity();
712
713        // Account for columns
714        size += self.columns.capacity() * std::mem::size_of::<DataColumn>();
715        for col in &self.columns {
716            size += col.name.capacity();
717        }
718
719        // Account for rows and their capacity
720        size += self.rows.capacity() * std::mem::size_of::<DataRow>();
721
722        // Account for actual data values
723        for row in &self.rows {
724            size += row.values.capacity() * std::mem::size_of::<DataValue>();
725            for value in &row.values {
726                match value {
727                    DataValue::String(s) => size += s.capacity(),
728                    DataValue::InternedString(_) => size += std::mem::size_of::<Arc<String>>(),
729                    DataValue::DateTime(s) => size += s.capacity(),
730                    _ => {} // Other types are inline
731                }
732            }
733        }
734
735        // Account for metadata
736        size += self.metadata.capacity() * std::mem::size_of::<(String, String)>();
737        for (k, v) in &self.metadata {
738            size += k.capacity() + v.capacity();
739        }
740
741        size
742    }
743}
744
745/// V46: Helper function to convert JSON value to DataValue
746fn json_value_to_data_value(json: &JsonValue) -> DataValue {
747    match json {
748        JsonValue::Null => DataValue::Null,
749        JsonValue::Bool(b) => DataValue::Boolean(*b),
750        JsonValue::Number(n) => {
751            if let Some(i) = n.as_i64() {
752                DataValue::Integer(i)
753            } else if let Some(f) = n.as_f64() {
754                DataValue::Float(f)
755            } else {
756                DataValue::String(n.to_string())
757            }
758        }
759        JsonValue::String(s) => {
760            // Try to detect if it's a date/time
761            if s.contains('-') && s.len() >= 8 && s.len() <= 30 {
762                // Simple heuristic for dates
763                DataValue::DateTime(s.clone())
764            } else {
765                DataValue::String(s.clone())
766            }
767        }
768        JsonValue::Array(_) | JsonValue::Object(_) => {
769            // Store complex types as JSON string
770            DataValue::String(json.to_string())
771        }
772    }
773}
774
775/// Statistics about a DataTable
776#[derive(Debug, Clone)]
777pub struct DataTableStats {
778    pub row_count: usize,
779    pub column_count: usize,
780    pub memory_size: usize,
781    pub null_count: usize,
782}
783
784/// Implementation of DataProvider for DataTable
785/// This allows DataTable to be used wherever DataProvider trait is expected
786impl DataProvider for DataTable {
787    fn get_row(&self, index: usize) -> Option<Vec<String>> {
788        self.rows
789            .get(index)
790            .map(|row| row.values.iter().map(|v| v.to_string_optimized()).collect())
791    }
792
793    fn get_column_names(&self) -> Vec<String> {
794        self.column_names()
795    }
796
797    fn get_row_count(&self) -> usize {
798        self.row_count()
799    }
800
801    fn get_column_count(&self) -> usize {
802        self.column_count()
803    }
804}
805
806#[cfg(test)]
807mod tests {
808    use super::*;
809
810    #[test]
811    fn test_data_type_inference() {
812        assert_eq!(DataType::infer_from_string("123"), DataType::Integer);
813        assert_eq!(DataType::infer_from_string("123.45"), DataType::Float);
814        assert_eq!(DataType::infer_from_string("true"), DataType::Boolean);
815        assert_eq!(DataType::infer_from_string("hello"), DataType::String);
816        assert_eq!(DataType::infer_from_string(""), DataType::Null);
817        assert_eq!(
818            DataType::infer_from_string("2024-01-01"),
819            DataType::DateTime
820        );
821    }
822
823    #[test]
824    fn test_datatable_creation() {
825        let mut table = DataTable::new("test");
826
827        table.add_column(DataColumn::new("id").with_type(DataType::Integer));
828        table.add_column(DataColumn::new("name").with_type(DataType::String));
829        table.add_column(DataColumn::new("active").with_type(DataType::Boolean));
830
831        assert_eq!(table.column_count(), 3);
832        assert_eq!(table.row_count(), 0);
833
834        let row = DataRow::new(vec![
835            DataValue::Integer(1),
836            DataValue::String("Alice".to_string()),
837            DataValue::Boolean(true),
838        ]);
839
840        table.add_row(row).unwrap();
841        assert_eq!(table.row_count(), 1);
842
843        let value = table.get_value_by_name(0, "name").unwrap();
844        assert_eq!(value.to_string(), "Alice");
845    }
846
847    #[test]
848    fn test_type_inference() {
849        let mut table = DataTable::new("test");
850
851        // Add columns without types
852        table.add_column(DataColumn::new("mixed"));
853
854        // Add rows with different types
855        table
856            .add_row(DataRow::new(vec![DataValue::Integer(1)]))
857            .unwrap();
858        table
859            .add_row(DataRow::new(vec![DataValue::Float(2.5)]))
860            .unwrap();
861        table.add_row(DataRow::new(vec![DataValue::Null])).unwrap();
862
863        table.infer_column_types();
864
865        // Should infer Float since we have both Integer and Float
866        assert_eq!(table.columns[0].data_type, DataType::Float);
867        assert_eq!(table.columns[0].null_count, 1);
868        assert!(table.columns[0].nullable);
869    }
870
871    #[test]
872    fn test_from_query_response() {
873        use crate::api_client::{QueryInfo, QueryResponse};
874        use serde_json::json;
875
876        let response = QueryResponse {
877            query: QueryInfo {
878                select: vec!["id".to_string(), "name".to_string(), "age".to_string()],
879                where_clause: None,
880                order_by: None,
881            },
882            data: vec![
883                json!({
884                    "id": 1,
885                    "name": "Alice",
886                    "age": 30
887                }),
888                json!({
889                    "id": 2,
890                    "name": "Bob",
891                    "age": 25
892                }),
893                json!({
894                    "id": 3,
895                    "name": "Carol",
896                    "age": null
897                }),
898            ],
899            count: 3,
900            source: Some("test.csv".to_string()),
901            table: Some("test".to_string()),
902            cached: Some(false),
903        };
904
905        let table = DataTable::from_query_response(&response, "test").unwrap();
906
907        assert_eq!(table.name, "test");
908        assert_eq!(table.row_count(), 3);
909        assert_eq!(table.column_count(), 3);
910
911        // Check column names
912        let col_names = table.column_names();
913        assert!(col_names.contains(&"id".to_string()));
914        assert!(col_names.contains(&"name".to_string()));
915        assert!(col_names.contains(&"age".to_string()));
916
917        // Check metadata
918        assert_eq!(table.metadata.get("source"), Some(&"test.csv".to_string()));
919        assert_eq!(table.metadata.get("cached"), Some(&"false".to_string()));
920
921        // Check first row values
922        assert_eq!(
923            table.get_value_by_name(0, "id"),
924            Some(&DataValue::Integer(1))
925        );
926        assert_eq!(
927            table.get_value_by_name(0, "name"),
928            Some(&DataValue::String("Alice".to_string()))
929        );
930        assert_eq!(
931            table.get_value_by_name(0, "age"),
932            Some(&DataValue::Integer(30))
933        );
934
935        // Check null handling
936        assert_eq!(table.get_value_by_name(2, "age"), Some(&DataValue::Null));
937    }
938}
sql_cli/data/datatable.rs

sql_cli/data/
datatable.rs
