pgx_utils/sql_entity_graph/extension_sql/
entity.rs

1/*
2Portions Copyright 2019-2021 ZomboDB, LLC.
3Portions Copyright 2021-2022 Technology Concepts & Design, Inc. <support@tcdi.com>
4
5All rights reserved.
6
7Use of this source code is governed by the MIT license that can be found in the LICENSE file.
8*/
9/*!
10
11`pgx::extension_sql!()` related entities for Rust to SQL translation
12
13> Like all of the [`sql_entity_graph`][crate::sql_entity_graph] APIs, this is considered **internal**
14to the `pgx` framework and very subject to change between versions. While you may use this, please do it with caution.
15
16
17*/
18use crate::sql_entity_graph::extension_sql::SqlDeclared;
19use crate::sql_entity_graph::pgx_sql::PgxSql;
20use crate::sql_entity_graph::positioning_ref::PositioningRef;
21use crate::sql_entity_graph::to_sql::ToSql;
22use crate::sql_entity_graph::{SqlGraphEntity, SqlGraphIdentifier};
23
24use std::fmt::Display;
25
26/// The output of a [`ExtensionSql`](crate::sql_entity_graph::ExtensionSql) from `quote::ToTokens::to_tokens`.
27#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
28pub struct ExtensionSqlEntity {
29    pub module_path: &'static str,
30    pub full_path: &'static str,
31    pub sql: &'static str,
32    pub file: &'static str,
33    pub line: u32,
34    pub name: &'static str,
35    pub bootstrap: bool,
36    pub finalize: bool,
37    pub requires: Vec<PositioningRef>,
38    pub creates: Vec<SqlDeclaredEntity>,
39}
40
41impl ExtensionSqlEntity {
42    pub fn has_sql_declared_entity(&self, identifier: &SqlDeclared) -> Option<&SqlDeclaredEntity> {
43        self.creates.iter().find(|created| created.has_sql_declared_entity(identifier))
44    }
45}
46
47impl From<ExtensionSqlEntity> for SqlGraphEntity {
48    fn from(val: ExtensionSqlEntity) -> Self {
49        SqlGraphEntity::CustomSql(val)
50    }
51}
52
53impl SqlGraphIdentifier for ExtensionSqlEntity {
54    fn dot_identifier(&self) -> String {
55        format!("sql {}", self.name)
56    }
57    fn rust_identifier(&self) -> String {
58        self.name.to_string()
59    }
60
61    fn file(&self) -> Option<&'static str> {
62        Some(self.file)
63    }
64
65    fn line(&self) -> Option<u32> {
66        Some(self.line)
67    }
68}
69
70impl ToSql for ExtensionSqlEntity {
71    #[tracing::instrument(level = "debug", skip(self, _context), fields(identifier = self.full_path))]
72    fn to_sql(&self, _context: &PgxSql) -> eyre::Result<String> {
73        let sql = format!(
74            "\n\
75                -- {file}:{line}\n\
76                {bootstrap}\
77                {creates}\
78                {requires}\
79                {finalize}\
80                {sql}\
81                ",
82            file = self.file,
83            line = self.line,
84            bootstrap = if self.bootstrap { "-- bootstrap\n" } else { "" },
85            creates = if !self.creates.is_empty() {
86                format!(
87                    "\
88                    -- creates:\n\
89                    {}\n\
90                ",
91                    self.creates
92                        .iter()
93                        .map(|i| format!("--   {}", i))
94                        .collect::<Vec<_>>()
95                        .join("\n")
96                ) + "\n"
97            } else {
98                "".to_string()
99            },
100            requires = if !self.requires.is_empty() {
101                format!(
102                    "\
103                   -- requires:\n\
104                    {}\n\
105                ",
106                    self.requires
107                        .iter()
108                        .map(|i| format!("--   {}", i))
109                        .collect::<Vec<_>>()
110                        .join("\n")
111                ) + "\n"
112            } else {
113                "".to_string()
114            },
115            finalize = if self.finalize { "-- finalize\n" } else { "" },
116            sql = self.sql,
117        );
118        tracing::trace!(%sql);
119        Ok(sql)
120    }
121}
122
123#[derive(Debug, Clone, Hash, PartialEq, Eq, Ord, PartialOrd)]
124pub struct SqlDeclaredEntityData {
125    sql: String,
126    name: String,
127    option: String,
128    vec: String,
129    vec_option: String,
130    option_vec: String,
131    option_vec_option: String,
132    array: String,
133    option_array: String,
134    varlena: String,
135    pg_box: Vec<String>,
136}
137#[derive(Debug, Clone, Hash, PartialEq, Eq, Ord, PartialOrd)]
138pub enum SqlDeclaredEntity {
139    Type(SqlDeclaredEntityData),
140    Enum(SqlDeclaredEntityData),
141    Function(SqlDeclaredEntityData),
142}
143
144impl Display for SqlDeclaredEntity {
145    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
146        match self {
147            SqlDeclaredEntity::Type(data) => {
148                f.write_str(&(String::from("Type(") + &data.name + ")"))
149            }
150            SqlDeclaredEntity::Enum(data) => {
151                f.write_str(&(String::from("Enum(") + &data.name + ")"))
152            }
153            SqlDeclaredEntity::Function(data) => {
154                f.write_str(&(String::from("Function ") + &data.name + ")"))
155            }
156        }
157    }
158}
159
160impl SqlDeclaredEntity {
161    pub fn build(variant: impl AsRef<str>, name: impl AsRef<str>) -> eyre::Result<Self> {
162        let name = name.as_ref();
163        let data = SqlDeclaredEntityData {
164            sql: name
165                .split("::")
166                .last()
167                .ok_or_else(|| eyre::eyre!("Did not get SQL for `{}`", name))?
168                .to_string(),
169            name: name.to_string(),
170            option: format!("Option<{}>", name),
171            vec: format!("Vec<{}>", name),
172            vec_option: format!("Vec<Option<{}>>", name),
173            option_vec: format!("Option<Vec<{}>>", name),
174            option_vec_option: format!("Option<Vec<Option<{}>>", name),
175            array: format!("Array<{}>", name),
176            option_array: format!("Option<{}>", name),
177            varlena: format!("Varlena<{}>", name),
178            pg_box: vec![
179                format!("pgx::pgbox::PgBox<{}>", name),
180                format!("pgx::pgbox::PgBox<{}, pgx::pgbox::AllocatedByRust>", name),
181                format!("pgx::pgbox::PgBox<{}, pgx::pgbox::AllocatedByPostgres>", name),
182            ],
183        };
184        let retval = match variant.as_ref() {
185            "Type" => Self::Type(data),
186            "Enum" => Self::Enum(data),
187            "Function" => Self::Function(data),
188            _ => {
189                return Err(eyre::eyre!(
190                    "Can only declare `Type(Ident)`, `Enum(Ident)` or `Function(Ident)`"
191                ))
192            }
193        };
194        Ok(retval)
195    }
196    pub fn sql(&self) -> String {
197        match self {
198            SqlDeclaredEntity::Type(data) => data.sql.clone(),
199            SqlDeclaredEntity::Enum(data) => data.sql.clone(),
200            SqlDeclaredEntity::Function(data) => data.sql.clone(),
201        }
202    }
203
204    pub fn has_sql_declared_entity(&self, identifier: &SqlDeclared) -> bool {
205        match (&identifier, &self) {
206            (SqlDeclared::Type(identifier_name), &SqlDeclaredEntity::Type(data))
207            | (SqlDeclared::Enum(identifier_name), &SqlDeclaredEntity::Enum(data))
208            | (SqlDeclared::Function(identifier_name), &SqlDeclaredEntity::Function(data)) => {
209                let matches = |identifier_name: &str| {
210                    identifier_name == &data.name
211                        || identifier_name == &data.option
212                        || identifier_name == &data.vec
213                        || identifier_name == &data.vec_option
214                        || identifier_name == &data.option_vec
215                        || identifier_name == &data.option_vec_option
216                        || identifier_name == &data.array
217                        || identifier_name == &data.option_array
218                        || identifier_name == &data.varlena
219                };
220                if matches(&*identifier_name) || data.pg_box.contains(&identifier_name) {
221                    return true;
222                }
223                // there are cases where the identifier is
224                // `core::option::Option<Foo>` while the data stores
225                // `Option<Foo>` check again for this
226                let generics_start = match identifier_name.find('<') {
227                    None => return false,
228                    Some(idx) => idx,
229                };
230                let qualification_end = match identifier_name[..generics_start].rfind("::") {
231                    None => return false,
232                    Some(idx) => idx,
233                };
234                matches(&identifier_name[qualification_end + 2..])
235            }
236            _ => false,
237        }
238    }
239}