datafusion_table_providers/sql/arrow_sql_gen/postgres/
builder.rs

1use datafusion::arrow::datatypes::Fields;
2use sea_query::{Alias, ColumnDef, PostgresQueryBuilder, TableBuilder};
3
4use crate::sql::arrow_sql_gen::statement::map_data_type_to_column_type;
5
6pub struct TypeBuilder {
7    name: String,
8    columns: Vec<ColumnDef>,
9}
10
11impl TypeBuilder {
12    #[must_use]
13    pub fn new(name: String, fields: &Fields) -> Self {
14        Self {
15            name,
16            columns: fields_to_simple_column_defs(fields),
17        }
18    }
19
20    #[must_use]
21    pub fn build(self) -> String {
22        let pg_builder = PostgresQueryBuilder;
23
24        let mut sql = String::new();
25
26        // Postgres doesn't natively support a CREATE TYPE IF NOT EXISTS statement,
27        // so we'll wrap it in a DO block with a conditional check.
28
29        sql.push_str(&format!(
30            "
31        DO $$ 
32        BEGIN
33            IF NOT EXISTS (
34                SELECT 1
35                FROM pg_type t
36                WHERE t.typname = '{}'
37            ) THEN
38                ",
39            self.name
40        ));
41
42        sql.push_str(&format!("CREATE TYPE {} AS (", self.name));
43
44        let mut first = true;
45
46        for column_def in &self.columns {
47            if !first {
48                sql.push_str(", ");
49            }
50
51            pg_builder.prepare_column_def(column_def, &mut sql);
52            first = false;
53        }
54
55        sql.push_str(" );");
56
57        sql.push_str(
58            "
59            END IF;
60        END $$;
61        ",
62        );
63
64        sql
65    }
66}
67
68/// Convert a `Fields` struct into a vector of `ColumnDef` without any constraints or other column specs.
69fn fields_to_simple_column_defs(fields: &Fields) -> Vec<ColumnDef> {
70    let mut column_defs = Vec::new();
71    for field in fields {
72        let column_type = map_data_type_to_column_type(field.data_type());
73        let column_def = ColumnDef::new_with_type(Alias::new(field.name()), column_type);
74
75        column_defs.push(column_def);
76    }
77
78    column_defs
79}
80
81#[cfg(test)]
82mod tests {
83    use datafusion::arrow::datatypes::{DataType, Field, Schema};
84
85    use super::*;
86
87    #[test]
88    fn test_type_builder() {
89        let fields = vec![
90            Field::new("id", DataType::Int32, false),
91            Field::new("name", DataType::Utf8, false),
92        ];
93        let schema = Schema::new(fields);
94
95        let type_builder = TypeBuilder::new("person".to_string(), schema.fields());
96        let sql = type_builder.build();
97
98        assert_eq!(
99            sql,
100            r#"
101        DO $$ 
102        BEGIN
103            IF NOT EXISTS (
104                SELECT 1
105                FROM pg_type t
106                WHERE t.typname = 'person'
107            ) THEN
108                CREATE TYPE person AS ("id" integer, "name" text );
109            END IF;
110        END $$;
111        "#
112        );
113    }
114}
datafusion_table_providers/sql/arrow_sql_gen/postgres/builder.rs

datafusion_table_providers/sql/arrow_sql_gen/postgres/
builder.rs
