pgrx_sql_entity_graph/pg_extern/entity/
mod.rs

1//LICENSE Portions Copyright 2019-2021 ZomboDB, LLC.
2//LICENSE
3//LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc.
4//LICENSE
5//LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. <contact@pgcentral.org>
6//LICENSE
7//LICENSE All rights reserved.
8//LICENSE
9//LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file.
10/*!
11
12`#[pg_extern]` related entities for Rust to SQL translation
13
14> Like all of the [`sql_entity_graph`][crate] APIs, this is considered **internal**
15> to the `pgrx` framework and very subject to change between versions. While you may use this, please do it with caution.
16
17*/
18mod argument;
19mod cast;
20mod operator;
21mod returning;
22
23pub use argument::PgExternArgumentEntity;
24pub use cast::PgCastEntity;
25pub use operator::PgOperatorEntity;
26pub use returning::{PgExternReturnEntity, PgExternReturnEntityIteratedItem};
27
28use crate::fmt;
29use crate::metadata::{Returns, SqlMapping};
30use crate::pgrx_sql::PgrxSql;
31use crate::to_sql::entity::ToSqlConfigEntity;
32use crate::to_sql::ToSql;
33use crate::{ExternArgs, SqlGraphEntity, SqlGraphIdentifier, TypeMatch};
34
35use eyre::{eyre, WrapErr};
36
37/// The output of a [`PgExtern`](crate::pg_extern::PgExtern) from `quote::ToTokens::to_tokens`.
38#[derive(Debug, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
39pub struct PgExternEntity {
40    pub name: &'static str,
41    pub unaliased_name: &'static str,
42    pub module_path: &'static str,
43    pub full_path: &'static str,
44    pub metadata: crate::metadata::FunctionMetadataEntity,
45    pub fn_args: Vec<PgExternArgumentEntity>,
46    pub fn_return: PgExternReturnEntity,
47    pub schema: Option<&'static str>,
48    pub file: &'static str,
49    pub line: u32,
50    pub extern_attrs: Vec<ExternArgs>,
51    pub search_path: Option<Vec<&'static str>>,
52    pub operator: Option<PgOperatorEntity>,
53    pub cast: Option<PgCastEntity>,
54    pub to_sql_config: ToSqlConfigEntity,
55}
56
57impl From<PgExternEntity> for SqlGraphEntity {
58    fn from(val: PgExternEntity) -> Self {
59        SqlGraphEntity::Function(val)
60    }
61}
62
63impl SqlGraphIdentifier for PgExternEntity {
64    fn dot_identifier(&self) -> String {
65        format!("fn {}", self.name)
66    }
67    fn rust_identifier(&self) -> String {
68        self.full_path.to_string()
69    }
70
71    fn file(&self) -> Option<&'static str> {
72        Some(self.file)
73    }
74
75    fn line(&self) -> Option<u32> {
76        Some(self.line)
77    }
78}
79
80impl PgExternEntity {
81    fn sql_name(&self, context: &PgrxSql) -> String {
82        let self_index = context.externs[self];
83        let schema = self
84            .schema
85            .map(|schema| format!("{schema}."))
86            .unwrap_or_else(|| context.schema_prefix_for(&self_index));
87
88        format!("{schema}\"{}\"", self.name)
89    }
90}
91
92impl ToSql for PgExternEntity {
93    fn to_sql(&self, context: &PgrxSql) -> eyre::Result<String> {
94        let self_index = context.externs[self];
95        let mut extern_attrs = self.extern_attrs.clone();
96        // if we already have a STRICT marker we do not need to add it
97        // presume we can upgrade, then disprove it
98        let mut strict_upgrade = !extern_attrs.iter().any(|i| i == &ExternArgs::Strict);
99        if strict_upgrade {
100            // It may be possible to infer a `STRICT` marker though.
101            // But we can only do that if the user hasn't used `Option<T>` or `pgrx::Internal`
102            for arg in &self.metadata.arguments {
103                if arg.optional {
104                    strict_upgrade = false;
105                }
106            }
107        }
108
109        if strict_upgrade {
110            extern_attrs.push(ExternArgs::Strict);
111        }
112        extern_attrs.sort();
113        extern_attrs.dedup();
114
115        let module_pathname = &context.get_module_pathname();
116        let schema = self
117            .schema
118            .map(|schema| format!("{schema}."))
119            .unwrap_or_else(|| context.schema_prefix_for(&self_index));
120        let arguments = if !self.fn_args.is_empty() {
121            let mut args = Vec::new();
122            let metadata_without_arg_skips = &self
123                .metadata
124                .arguments
125                .iter()
126                .filter(|v| v.argument_sql != Ok(SqlMapping::Skip))
127                .collect::<Vec<_>>();
128            for (idx, arg) in self.fn_args.iter().enumerate() {
129                let graph_index = context
130                    .graph
131                    .neighbors_undirected(self_index)
132                    .find(|neighbor| context.graph[*neighbor].type_matches(&arg.used_ty))
133                    .ok_or_else(|| eyre!("Could not find arg type in graph. Got: {:?}", arg))?;
134                let needs_comma = idx < (metadata_without_arg_skips.len().saturating_sub(1));
135                let metadata_argument = &self.metadata.arguments[idx];
136                match metadata_argument.argument_sql {
137                    Ok(SqlMapping::As(ref argument_sql)) => {
138                        let buf = format!("\
139                                            \t\"{pattern}\" {variadic}{schema_prefix}{sql_type}{default}{maybe_comma}/* {type_name} */\
140                                        ",
141                                          pattern = arg.pattern,
142                                          schema_prefix = context.schema_prefix_for(&graph_index),
143                                          // First try to match on [`TypeId`] since it's most reliable.
144                                          sql_type = argument_sql,
145                                          default = if let Some(def) = arg.used_ty.default { format!(" DEFAULT {def}") } else { String::from("") },
146                                          variadic = if metadata_argument.variadic { "VARIADIC " } else { "" },
147                                          maybe_comma = if needs_comma { ", " } else { " " },
148                                          type_name = metadata_argument.type_name,
149                        );
150                        args.push(buf);
151                    }
152                    Ok(SqlMapping::Composite { array_brackets }) => {
153                        let sql = self.fn_args[idx]
154                            .used_ty
155                            .composite_type
156                            .map(|v| fmt::with_array_brackets(v.into(), array_brackets))
157                            .ok_or_else(|| {
158                                eyre!(
159                                    "Macro expansion time suggested a composite_type!() in return"
160                                )
161                            })?;
162                        let buf = format!("\
163                            \t\"{pattern}\" {variadic}{schema_prefix}{sql_type}{default}{maybe_comma}/* {type_name} */\
164                        ",
165                                          pattern = arg.pattern,
166                                          schema_prefix = context.schema_prefix_for(&graph_index),
167                                          // First try to match on [`TypeId`] since it's most reliable.
168                                          sql_type = sql,
169                                          default = if let Some(def) = arg.used_ty.default { format!(" DEFAULT {def}") } else { String::from("") },
170                                          variadic = if metadata_argument.variadic { "VARIADIC " } else { "" },
171                                          maybe_comma = if needs_comma { ", " } else { " " },
172                                          type_name = metadata_argument.type_name,
173                        );
174                        args.push(buf);
175                    }
176                    Ok(SqlMapping::Skip) => (),
177                    Err(err) => return Err(err).wrap_err("While mapping argument"),
178                }
179            }
180            String::from("\n") + &args.join("\n") + "\n"
181        } else {
182            Default::default()
183        };
184
185        let returns = match &self.fn_return {
186            PgExternReturnEntity::None => String::from("RETURNS void"),
187            PgExternReturnEntity::Type { ty } => {
188                let graph_index = context
189                    .graph
190                    .neighbors_undirected(self_index)
191                    .find(|neighbor| context.graph[*neighbor].type_matches(ty))
192                    .ok_or_else(|| eyre!("Could not find return type in graph."))?;
193                let metadata_retval = self.metadata.retval.clone();
194                let sql_type = match metadata_retval.return_sql {
195                    Ok(Returns::One(SqlMapping::As(ref sql))) => sql.clone(),
196                    Ok(Returns::One(SqlMapping::Composite { array_brackets })) => fmt::with_array_brackets(ty.composite_type.unwrap().into(), array_brackets),
197                    Ok(other) => return Err(eyre!("Got non-plain mapped/composite return variant SQL in what macro-expansion thought was a type, got: {other:?}")),
198                    Err(err) => return Err(err).wrap_err("Error mapping return SQL")
199                };
200                format!(
201                    "RETURNS {schema_prefix}{sql_type} /* {full_path} */",
202                    schema_prefix = context.schema_prefix_for(&graph_index),
203                    full_path = ty.full_path
204                )
205            }
206            PgExternReturnEntity::SetOf { ty, .. } => {
207                let graph_index = context
208                    .graph
209                    .neighbors_undirected(self_index)
210                    .find(|neighbor| context.graph[*neighbor].type_matches(ty))
211                    .ok_or_else(|| eyre!("Could not find return type in graph."))?;
212                let metadata_retval = self.metadata.retval.clone();
213                let sql_type = match metadata_retval.return_sql {
214                    Ok(Returns::SetOf(SqlMapping::As(ref sql))) => sql.clone(),
215                    Ok(Returns::SetOf(SqlMapping::Composite { array_brackets })) => fmt::with_array_brackets(ty.composite_type.unwrap().into(), array_brackets),
216                    Ok(_other) => return Err(eyre!("Got non-setof mapped/composite return variant SQL in what macro-expansion thought was a setof")),
217                    Err(err) => return Err(err).wrap_err("Error mapping return SQL"),
218                };
219                format!(
220                    "RETURNS SETOF {schema_prefix}{sql_type} /* {full_path} */",
221                    schema_prefix = context.schema_prefix_for(&graph_index),
222                    full_path = ty.full_path
223                )
224            }
225            PgExternReturnEntity::Iterated { tys: table_items, .. } => {
226                let mut items = String::new();
227                let metadata_retval = self.metadata.retval.clone();
228                let metadata_retval_sqls: Vec<String> = match metadata_retval.return_sql {
229                    Ok(Returns::Table(variants)) => {
230                        variants.iter().enumerate().map(|(idx, variant)| {
231                            match variant {
232                                SqlMapping::As(sql) => sql.clone(),
233                                SqlMapping::Composite { array_brackets } => {
234                                    let composite = table_items[idx].ty.composite_type.unwrap();
235                                    fmt::with_array_brackets(composite.into(), *array_brackets)
236                                }
237                                SqlMapping::Skip => todo!(),
238                            }
239                        }).collect()
240                    }
241                    Ok(_other) => return Err(eyre!("Got non-table return variant SQL in what macro-expansion thought was a table")),
242                    Err(err) => return Err(err).wrap_err("Error mapping return SQL"),
243                };
244
245                for (idx, returning::PgExternReturnEntityIteratedItem { ty, name: col_name }) in
246                    table_items.iter().enumerate()
247                {
248                    let graph_index =
249                        context.graph.neighbors_undirected(self_index).find(|neighbor| {
250                            context.graph[*neighbor].id_or_name_matches(&ty.ty_id, ty.ty_source)
251                        });
252
253                    let needs_comma = idx < (table_items.len() - 1);
254                    let item = format!(
255                        "\n\t{col_name} {schema_prefix}{ty_resolved}{needs_comma} /* {ty_name} */",
256                        col_name = col_name.expect(
257                            "An iterator of tuples should have `named!()` macro declarations."
258                        ),
259                        schema_prefix = if let Some(graph_index) = graph_index {
260                            context.schema_prefix_for(&graph_index)
261                        } else {
262                            "".into()
263                        },
264                        ty_resolved = metadata_retval_sqls[idx],
265                        needs_comma = if needs_comma { ", " } else { " " },
266                        ty_name = ty.full_path
267                    );
268                    items.push_str(&item);
269                }
270                format!("RETURNS TABLE ({items}\n)")
271            }
272            PgExternReturnEntity::Trigger => String::from("RETURNS trigger"),
273        };
274        let PgExternEntity { name, module_path, file, line, .. } = self;
275
276        let fn_sql = format!(
277            "\
278                CREATE {or_replace} FUNCTION {schema}\"{name}\"({arguments}) {returns}\n\
279                {extern_attrs}\
280                {search_path}\
281                LANGUAGE c /* Rust */\n\
282                AS '{module_pathname}', '{unaliased_name}_wrapper';\
283            ",
284            or_replace =
285                if extern_attrs.contains(&ExternArgs::CreateOrReplace) { "OR REPLACE" } else { "" },
286            search_path = if let Some(search_path) = &self.search_path {
287                let retval = format!("SET search_path TO {}", search_path.join(", "));
288                retval + "\n"
289            } else {
290                Default::default()
291            },
292            extern_attrs = if extern_attrs.is_empty() {
293                String::default()
294            } else {
295                let mut retval = extern_attrs
296                    .iter()
297                    .filter(|attr| **attr != ExternArgs::CreateOrReplace)
298                    .map(|attr| {
299                        if matches!(attr, ExternArgs::Support(..)) {
300                            let support_fn_name = attr.to_string();
301
302                            let support_fn_name =
303                            if let Some(entity) = context.find_matching_fn(&support_fn_name) {
304                                entity.sql_name(context)
305                            } else {
306                                panic!("cannot locate SUPPORT function `{support_fn_name}` attached to function `{}`", self.full_path)
307                            };
308
309                            format!("SUPPORT {support_fn_name}")
310                        } else {
311                            attr.to_string().to_uppercase()
312                        }
313                    })
314                    .collect::<Vec<_>>()
315                    .join(" ");
316                retval.push('\n');
317                retval
318            },
319            unaliased_name = self.unaliased_name,
320        );
321
322        let requires = {
323            let requires_attrs = self
324                .extern_attrs
325                .iter()
326                .filter_map(|x| match x {
327                    ExternArgs::Requires(requirements) => Some(requirements.clone()),
328                    ExternArgs::Support(support_fn) => Some(vec![support_fn.clone()]),
329                    _ => None,
330                })
331                .flatten()
332                .collect::<Vec<_>>();
333
334            if !requires_attrs.is_empty() {
335                format!(
336                    "-- requires:\n{}\n",
337                    requires_attrs
338                        .iter()
339                        .map(|i| format!("--   {i}"))
340                        .collect::<Vec<_>>()
341                        .join("\n")
342                )
343            } else {
344                "".to_string()
345            }
346        };
347
348        let mut ext_sql = format!(
349            "\n\
350            -- {file}:{line}\n\
351            -- {module_path}::{name}\n\
352            {requires}\
353            {fn_sql}"
354        );
355
356        if let Some(op) = &self.operator {
357            let mut optionals = vec![];
358            if let Some(it) = op.commutator {
359                optionals.push(format!("\tCOMMUTATOR = {it}"));
360            };
361            if let Some(it) = op.negator {
362                optionals.push(format!("\tNEGATOR = {it}"));
363            };
364            if let Some(it) = op.restrict {
365                optionals.push(format!("\tRESTRICT = {it}"));
366            };
367            if let Some(it) = op.join {
368                optionals.push(format!("\tJOIN = {it}"));
369            };
370            if op.hashes {
371                optionals.push(String::from("\tHASHES"));
372            };
373            if op.merges {
374                optionals.push(String::from("\tMERGES"));
375            };
376
377            let left_arg =
378                self.metadata.arguments.first().ok_or_else(|| {
379                    eyre!("Did not find `left_arg` for operator `{}`.", self.name)
380                })?;
381            let left_fn_arg = self
382                .fn_args
383                .first()
384                .ok_or_else(|| eyre!("Did not find `left_arg` for operator `{}`.", self.name))?;
385            let left_arg_graph_index = context
386                .graph
387                .neighbors_undirected(self_index)
388                .find(|neighbor| {
389                    context.graph[*neighbor]
390                        .id_or_name_matches(&left_fn_arg.used_ty.ty_id, left_arg.type_name)
391                })
392                .ok_or_else(|| {
393                    eyre!("Could not find left arg type in graph. Got: {:?}", left_arg)
394                })?;
395            let left_arg_sql = match left_arg.argument_sql {
396                Ok(SqlMapping::As(ref sql)) => sql.clone(),
397                Ok(SqlMapping::Composite { array_brackets }) => {
398                    if array_brackets {
399                        let composite_type = self.fn_args[0].used_ty.composite_type
400                            .ok_or(eyre!("Found a composite type but macro expansion time did not reveal a name, use `pgrx::composite_type!()`"))?;
401                        format!("{composite_type}[]")
402                    } else {
403                        self.fn_args[0].used_ty.composite_type
404                            .ok_or(eyre!("Found a composite type but macro expansion time did not reveal a name, use `pgrx::composite_type!()`"))?.to_string()
405                    }
406                }
407                Ok(SqlMapping::Skip) => {
408                    return Err(eyre!(
409                        "Found an skipped SQL type in an operator, this is not valid"
410                    ))
411                }
412                Err(err) => return Err(err.into()),
413            };
414
415            let right_arg =
416                self.metadata.arguments.get(1).ok_or_else(|| {
417                    eyre!("Did not find `left_arg` for operator `{}`.", self.name)
418                })?;
419            let right_fn_arg = self
420                .fn_args
421                .get(1)
422                .ok_or_else(|| eyre!("Did not find `left_arg` for operator `{}`.", self.name))?;
423            let right_arg_graph_index = context
424                .graph
425                .neighbors_undirected(self_index)
426                .find(|neighbor| {
427                    context.graph[*neighbor]
428                        .id_or_name_matches(&right_fn_arg.used_ty.ty_id, right_arg.type_name)
429                })
430                .ok_or_else(|| {
431                    eyre!("Could not find right arg type in graph. Got: {:?}", right_arg)
432                })?;
433            let right_arg_sql = match right_arg.argument_sql {
434                Ok(SqlMapping::As(ref sql)) => sql.clone(),
435                Ok(SqlMapping::Composite { array_brackets }) => {
436                    if array_brackets {
437                        let composite_type = self.fn_args[1].used_ty.composite_type
438                            .ok_or(eyre!("Found a composite type but macro expansion time did not reveal a name, use `pgrx::composite_type!()`"))?;
439                        format!("{composite_type}[]")
440                    } else {
441                        self.fn_args[0].used_ty.composite_type
442                            .ok_or(eyre!("Found a composite type but macro expansion time did not reveal a name, use `pgrx::composite_type!()`"))?.to_string()
443                    }
444                }
445                Ok(SqlMapping::Skip) => {
446                    return Err(eyre!(
447                        "Found an skipped SQL type in an operator, this is not valid"
448                    ))
449                }
450                Err(err) => return Err(err.into()),
451            };
452
453            let schema = self
454                .schema
455                .map(|schema| format!("{schema}."))
456                .unwrap_or_else(|| context.schema_prefix_for(&self_index));
457
458            let operator_sql = format!("\n\n\
459                                                    -- {file}:{line}\n\
460                                                    -- {module_path}::{name}\n\
461                                                    CREATE OPERATOR {schema}{opname} (\n\
462                                                        \tPROCEDURE={schema}\"{name}\",\n\
463                                                        \tLEFTARG={schema_prefix_left}{left_arg_sql}, /* {left_name} */\n\
464                                                        \tRIGHTARG={schema_prefix_right}{right_arg_sql}{maybe_comma} /* {right_name} */\n\
465                                                        {optionals}\
466                                                    );\
467                                                    ",
468                                       opname = op.opname.unwrap(),
469                                       left_name = left_arg.type_name,
470                                       right_name = right_arg.type_name,
471                                       schema_prefix_left = context.schema_prefix_for(&left_arg_graph_index),
472                                       schema_prefix_right = context.schema_prefix_for(&right_arg_graph_index),
473                                       maybe_comma = if !optionals.is_empty() { "," } else { "" },
474                                       optionals = if !optionals.is_empty() { optionals.join(",\n") + "\n" } else { "".to_string() },
475            );
476            ext_sql += &operator_sql
477        };
478        if let Some(cast) = &self.cast {
479            let target_arg = &self.metadata.retval;
480            let target_fn_arg = &self.fn_return;
481            let target_arg_graph_index = context
482                .graph
483                .neighbors_undirected(self_index)
484                .find(|neighbor| match (&context.graph[*neighbor], target_fn_arg) {
485                    (SqlGraphEntity::Type(ty), PgExternReturnEntity::Type { ty: rty }) => {
486                        ty.id_matches(&rty.ty_id)
487                    }
488                    (SqlGraphEntity::Enum(en), PgExternReturnEntity::Type { ty: rty }) => {
489                        en.id_matches(&rty.ty_id)
490                    }
491                    (SqlGraphEntity::BuiltinType(defined), _) => defined == target_arg.type_name,
492                    _ => false,
493                })
494                .ok_or_else(|| {
495                    eyre!("Could not find source type in graph. Got: {:?}", target_arg)
496                })?;
497            let target_arg_sql = match target_arg.argument_sql {
498                Ok(SqlMapping::As(ref sql)) => sql.clone(),
499                Ok(SqlMapping::Composite { array_brackets }) => {
500                    if array_brackets {
501                        let composite_type = self.fn_args[0].used_ty.composite_type
502                            .ok_or(eyre!("Found a composite type but macro expansion time did not reveal a name, use `pgrx::composite_type!()`"))?;
503                        format!("{composite_type}[]")
504                    } else {
505                        self.fn_args[0].used_ty.composite_type
506                            .ok_or(eyre!("Found a composite type but macro expansion time did not reveal a name, use `pgrx::composite_type!()`"))?.to_string()
507                    }
508                }
509                Ok(SqlMapping::Skip) => {
510                    return Err(eyre!("Found an skipped SQL type in a cast, this is not valid"))
511                }
512                Err(err) => return Err(err.into()),
513            };
514            let source_arg = self
515                .metadata
516                .arguments
517                .first()
518                .ok_or_else(|| eyre!("Did not find source type for cast `{}`.", self.name))?;
519            let source_fn_arg = self
520                .fn_args
521                .first()
522                .ok_or_else(|| eyre!("Did not find source type for cast `{}`.", self.name))?;
523            let source_arg_graph_index = context
524                .graph
525                .neighbors_undirected(self_index)
526                .find(|neighbor| {
527                    context.graph[*neighbor]
528                        .id_or_name_matches(&source_fn_arg.used_ty.ty_id, source_arg.type_name)
529                })
530                .ok_or_else(|| {
531                    eyre!("Could not find source type in graph. Got: {:?}", source_arg)
532                })?;
533            let source_arg_sql = match source_arg.argument_sql {
534                Ok(SqlMapping::As(ref sql)) => sql.clone(),
535                Ok(SqlMapping::Composite { array_brackets }) => {
536                    if array_brackets {
537                        let composite_type = self.fn_args[0].used_ty.composite_type
538                            .ok_or(eyre!("Found a composite type but macro expansion time did not reveal a name, use `pgrx::composite_type!()`"))?;
539                        format!("{composite_type}[]")
540                    } else {
541                        self.fn_args[0].used_ty.composite_type
542                            .ok_or(eyre!("Found a composite type but macro expansion time did not reveal a name, use `pgrx::composite_type!()`"))?.to_string()
543                    }
544                }
545                Ok(SqlMapping::Skip) => {
546                    return Err(eyre!("Found an skipped SQL type in a cast, this is not valid"))
547                }
548                Err(err) => return Err(err.into()),
549            };
550            let optional = match cast {
551                PgCastEntity::Default => String::from(""),
552                PgCastEntity::Assignment => String::from(" AS ASSIGNMENT"),
553                PgCastEntity::Implicit => String::from(" AS IMPLICIT"),
554            };
555
556            let cast_sql = format!("\n\n\
557                                                    -- {file}:{line}\n\
558                                                    -- {module_path}::{name}\n\
559                                                    CREATE CAST (\n\
560                                                        \t{schema_prefix_source}{source_arg_sql} /* {source_name} */\n\
561                                                        \tAS\n\
562                                                        \t{schema_prefix_target}{target_arg_sql} /* {target_name} */\n\
563                                                    )\n\
564                                                    WITH FUNCTION {function_name}{optional};\
565                                                    ",
566                                   file = self.file,
567                                   line = self.line,
568                                   name = self.name,
569                                   module_path = self.module_path,
570                                   schema_prefix_source = context.schema_prefix_for(&source_arg_graph_index),
571                                   source_name = source_arg.type_name,
572                                   schema_prefix_target = context.schema_prefix_for(&target_arg_graph_index),
573                                   target_name = target_arg.type_name,
574                                   function_name = self.name,
575            );
576            ext_sql += &cast_sql
577        };
578        Ok(ext_sql)
579    }
580}