pgx_utils/sql_entity_graph/
pgx_sql.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
11Rust to SQL mapping support.
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 eyre::eyre;
19use petgraph::dot::Dot;
20use petgraph::graph::NodeIndex;
21use petgraph::stable_graph::StableGraph;
22use std::any::TypeId;
23use std::collections::HashMap;
24use std::fmt::Debug;
25use std::path::Path;
26use tracing::instrument;
27
28use crate::sql_entity_graph::aggregate::entity::PgAggregateEntity;
29use crate::sql_entity_graph::control_file::ControlFile;
30use crate::sql_entity_graph::extension_sql::entity::{ExtensionSqlEntity, SqlDeclaredEntity};
31use crate::sql_entity_graph::extension_sql::SqlDeclared;
32use crate::sql_entity_graph::mapping::RustSourceOnlySqlMapping;
33use crate::sql_entity_graph::pg_extern::entity::PgExternEntity;
34use crate::sql_entity_graph::pg_trigger::entity::PgTriggerEntity;
35use crate::sql_entity_graph::positioning_ref::PositioningRef;
36use crate::sql_entity_graph::postgres_enum::entity::PostgresEnumEntity;
37use crate::sql_entity_graph::postgres_hash::entity::PostgresHashEntity;
38use crate::sql_entity_graph::postgres_ord::entity::PostgresOrdEntity;
39use crate::sql_entity_graph::postgres_type::entity::PostgresTypeEntity;
40use crate::sql_entity_graph::schema::entity::SchemaEntity;
41use crate::sql_entity_graph::to_sql::ToSql;
42use crate::sql_entity_graph::{SqlGraphEntity, SqlGraphIdentifier};
43
44use super::{PgExternReturnEntity, PgExternReturnEntityIteratedItem};
45
46#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord)]
47pub enum SqlGraphRelationship {
48    RequiredBy,
49    RequiredByArg,
50    RequiredByReturn,
51}
52
53#[derive(Debug, Clone)]
54pub struct RustToSqlMapping {
55    pub rust_source_to_sql: std::collections::HashSet<RustSourceOnlySqlMapping>,
56}
57
58/// A generator for SQL.
59///
60/// Consumes a base mapping of types (typically `pgx::DEFAULT_TYPEID_SQL_MAPPING`), a
61/// [`ControlFile`], and collections of each SQL entity.
62///
63/// During construction, a Directed Acyclic Graph is formed out the dependencies. For example,
64/// an item `detect_dog(x: &[u8]) -> animals::Dog` would have have a relationship with
65/// `animals::Dog`.
66///
67/// Typically, [`PgxSql`] types are constructed in a `pgx::pg_binary_magic!()` call in a binary
68/// out of entities collected during a `pgx::pg_module_magic!()` call in a library.
69#[derive(Debug, Clone)]
70pub struct PgxSql {
71    pub source_mappings: HashMap<String, RustSourceOnlySqlMapping>,
72    pub control: ControlFile,
73    pub graph: StableGraph<SqlGraphEntity, SqlGraphRelationship>,
74    pub graph_root: NodeIndex,
75    pub graph_bootstrap: Option<NodeIndex>,
76    pub graph_finalize: Option<NodeIndex>,
77    pub schemas: HashMap<SchemaEntity, NodeIndex>,
78    pub extension_sqls: HashMap<ExtensionSqlEntity, NodeIndex>,
79    pub externs: HashMap<PgExternEntity, NodeIndex>,
80    pub types: HashMap<PostgresTypeEntity, NodeIndex>,
81    pub builtin_types: HashMap<String, NodeIndex>,
82    pub enums: HashMap<PostgresEnumEntity, NodeIndex>,
83    pub ords: HashMap<PostgresOrdEntity, NodeIndex>,
84    pub hashes: HashMap<PostgresHashEntity, NodeIndex>,
85    pub aggregates: HashMap<PgAggregateEntity, NodeIndex>,
86    pub triggers: HashMap<PgTriggerEntity, NodeIndex>,
87    pub extension_name: String,
88    pub versioned_so: bool,
89}
90
91impl PgxSql {
92    #[instrument(level = "error", skip(sql_mappings, entities,))]
93    pub fn build(
94        sql_mappings: RustToSqlMapping,
95        entities: impl Iterator<Item = SqlGraphEntity>,
96        extension_name: String,
97        versioned_so: bool,
98    ) -> eyre::Result<Self> {
99        let RustToSqlMapping { rust_source_to_sql: source_mappings } = sql_mappings;
100
101        let mut graph = StableGraph::new();
102
103        let mut entities = entities.collect::<Vec<_>>();
104        entities.sort();
105        // Split up things into their specific types:
106        let mut control: Option<ControlFile> = None;
107        let mut schemas: Vec<SchemaEntity> = Vec::default();
108        let mut extension_sqls: Vec<ExtensionSqlEntity> = Vec::default();
109        let mut externs: Vec<PgExternEntity> = Vec::default();
110        let mut types: Vec<PostgresTypeEntity> = Vec::default();
111        let mut enums: Vec<PostgresEnumEntity> = Vec::default();
112        let mut ords: Vec<PostgresOrdEntity> = Vec::default();
113        let mut hashes: Vec<PostgresHashEntity> = Vec::default();
114        let mut aggregates: Vec<PgAggregateEntity> = Vec::default();
115        let mut triggers: Vec<PgTriggerEntity> = Vec::default();
116        for entity in entities {
117            match entity {
118                SqlGraphEntity::ExtensionRoot(input_control) => {
119                    control = Some(input_control);
120                }
121                SqlGraphEntity::Schema(input_schema) => {
122                    schemas.push(input_schema);
123                }
124                SqlGraphEntity::CustomSql(input_sql) => {
125                    extension_sqls.push(input_sql);
126                }
127                SqlGraphEntity::Function(input_function) => {
128                    externs.push(input_function);
129                }
130                SqlGraphEntity::Type(input_type) => {
131                    types.push(input_type);
132                }
133                SqlGraphEntity::BuiltinType(_) => (),
134                SqlGraphEntity::Enum(input_enum) => {
135                    enums.push(input_enum);
136                }
137                SqlGraphEntity::Ord(input_ord) => {
138                    ords.push(input_ord);
139                }
140                SqlGraphEntity::Hash(input_hash) => {
141                    hashes.push(input_hash);
142                }
143                SqlGraphEntity::Aggregate(input_aggregate) => {
144                    aggregates.push(input_aggregate);
145                }
146                SqlGraphEntity::Trigger(input_trigger) => {
147                    triggers.push(input_trigger);
148                }
149            }
150        }
151
152        let control: ControlFile = control.expect("No control file found");
153        let root = graph.add_node(SqlGraphEntity::ExtensionRoot(control.clone()));
154
155        // The initial build phase.
156        //
157        // Notably, we do not set non-root edges here. We do that in a second step. This is
158        // primarily because externs, types, operators, and the like tend to intertwine. If we tried
159        // to do it here, we'd find ourselves trying to create edges to non-existing entities.
160
161        // Both of these must be unique, so we can only hold one.
162        // Populate nodes, but don't build edges until we know if there is a bootstrap/finalize.
163        let (mapped_extension_sqls, bootstrap, finalize) =
164            initialize_extension_sqls(&mut graph, root, extension_sqls)?;
165        let mapped_schemas = initialize_schemas(&mut graph, bootstrap, finalize, schemas)?;
166        let mapped_enums = initialize_enums(&mut graph, root, bootstrap, finalize, enums)?;
167        let mapped_types = initialize_types(&mut graph, root, bootstrap, finalize, types)?;
168        let (mapped_externs, mut mapped_builtin_types) = initialize_externs(
169            &mut graph,
170            root,
171            bootstrap,
172            finalize,
173            externs,
174            &mapped_types,
175            &mapped_enums,
176        )?;
177        let mapped_ords = initialize_ords(&mut graph, root, bootstrap, finalize, ords)?;
178        let mapped_hashes = initialize_hashes(&mut graph, root, bootstrap, finalize, hashes)?;
179        let mapped_aggregates = initialize_aggregates(
180            &mut graph,
181            root,
182            bootstrap,
183            finalize,
184            aggregates,
185            &mut mapped_builtin_types,
186            &mapped_enums,
187            &mapped_types,
188        )?;
189        let mapped_triggers = initialize_triggers(&mut graph, root, bootstrap, finalize, triggers)?;
190
191        // Now we can circle back and build up the edge sets.
192        connect_schemas(&mut graph, &mapped_schemas, root);
193        connect_extension_sqls(
194            &mut graph,
195            &mapped_extension_sqls,
196            &mapped_schemas,
197            &mapped_types,
198            &mapped_enums,
199            &mapped_externs,
200            &mapped_triggers,
201        )?;
202        connect_enums(&mut graph, &mapped_enums, &mapped_schemas);
203        connect_types(&mut graph, &mapped_types, &mapped_schemas);
204        connect_externs(
205            &mut graph,
206            &mapped_externs,
207            &mapped_hashes,
208            &mapped_schemas,
209            &mapped_types,
210            &mapped_enums,
211            &mapped_builtin_types,
212            &mapped_extension_sqls,
213            &mapped_triggers,
214        )?;
215        connect_ords(
216            &mut graph,
217            &mapped_ords,
218            &mapped_schemas,
219            &mapped_types,
220            &mapped_enums,
221            &mapped_externs,
222        );
223        connect_hashes(
224            &mut graph,
225            &mapped_hashes,
226            &mapped_schemas,
227            &mapped_types,
228            &mapped_enums,
229            &mapped_externs,
230        );
231        connect_aggregates(
232            &mut graph,
233            &mapped_aggregates,
234            &mapped_schemas,
235            &mapped_types,
236            &mapped_enums,
237            &mapped_builtin_types,
238            &mapped_externs,
239        )?;
240        connect_triggers(&mut graph, &mapped_triggers, &mapped_schemas);
241
242        let this = Self {
243            source_mappings: source_mappings.into_iter().map(|x| (x.rust.clone(), x)).collect(),
244            control: control,
245            schemas: mapped_schemas,
246            extension_sqls: mapped_extension_sqls,
247            externs: mapped_externs,
248            types: mapped_types,
249            builtin_types: mapped_builtin_types,
250            enums: mapped_enums,
251            ords: mapped_ords,
252            hashes: mapped_hashes,
253            aggregates: mapped_aggregates,
254            triggers: mapped_triggers,
255            graph: graph,
256            graph_root: root,
257            graph_bootstrap: bootstrap,
258            graph_finalize: finalize,
259            extension_name: extension_name,
260            versioned_so,
261        };
262        Ok(this)
263    }
264
265    #[instrument(level = "error", skip(self))]
266    pub fn to_file(&self, file: impl AsRef<Path> + Debug) -> eyre::Result<()> {
267        use std::fs::{create_dir_all, File};
268        use std::io::Write;
269        use std::path::Path;
270        let generated = self.to_sql()?;
271        let path = Path::new(file.as_ref());
272
273        let parent = path.parent();
274        if let Some(parent) = parent {
275            create_dir_all(parent)?;
276        }
277        let mut out = File::create(path)?;
278        write!(out, "{}", generated)?;
279        Ok(())
280    }
281
282    #[instrument(level = "error", skip_all)]
283    pub fn write(&self, out: &mut impl std::io::Write) -> eyre::Result<()> {
284        let generated = self.to_sql()?;
285        if atty::is(atty::Stream::Stdout) {
286            self.write_highlighted(out, &generated)?;
287        } else {
288            write!(*out, "{}", generated)?;
289        }
290
291        Ok(())
292    }
293
294    #[cfg(not(feature = "syntax-highlighting"))]
295    fn write_highlighted(&self, out: &mut dyn std::io::Write, generated: &str) -> eyre::Result<()> {
296        write!(*out, "{}", generated)?;
297        Ok(())
298    }
299
300    #[cfg(feature = "syntax-highlighting")]
301    fn write_highlighted(&self, out: &mut dyn std::io::Write, generated: &str) -> eyre::Result<()> {
302        use eyre::WrapErr as _;
303        use owo_colors::{OwoColorize, XtermColors};
304        use syntect::easy::HighlightLines;
305        use syntect::highlighting::{Style, ThemeSet};
306        use syntect::parsing::SyntaxSet;
307        use syntect::util::LinesWithEndings;
308        let ps = SyntaxSet::load_defaults_newlines();
309        let theme_bytes = include_str!("../../assets/ansi.tmTheme").as_bytes();
310        let mut theme_reader = std::io::Cursor::new(theme_bytes);
311        let theme = ThemeSet::load_from_reader(&mut theme_reader)
312            .wrap_err("Couldn't parse theme for SQL highlighting, try piping to a file")?;
313
314        if let Some(syntax) = ps.find_syntax_by_extension("sql") {
315            let mut h = HighlightLines::new(syntax, &theme);
316            for line in LinesWithEndings::from(&generated) {
317                let ranges: Vec<(Style, &str)> = h.highlight_line(line, &ps)?;
318                // Concept from https://github.com/sharkdp/bat/blob/1b030dc03b906aa345f44b8266bffeea77d763fe/src/terminal.rs#L6
319                for (style, content) in ranges {
320                    if style.foreground.a == 0x01 {
321                        write!(*out, "{}", content)?;
322                    } else {
323                        write!(*out, "{}", content.color(XtermColors::from(style.foreground.r)))?;
324                    }
325                }
326                write!(*out, "\x1b[0m")?;
327            }
328        } else {
329            write!(*out, "{}", generated)?;
330        }
331        Ok(())
332    }
333
334    #[instrument(level = "error", err, skip(self))]
335    pub fn to_dot(&self, file: impl AsRef<Path> + Debug) -> eyre::Result<()> {
336        use std::fs::{create_dir_all, File};
337        use std::io::Write;
338        use std::path::Path;
339        let generated = Dot::with_attr_getters(
340            &self.graph,
341            &[petgraph::dot::Config::EdgeNoLabel, petgraph::dot::Config::NodeNoLabel],
342            &|_graph, edge| match edge.weight() {
343                SqlGraphRelationship::RequiredBy => format!(r#"color = "gray""#),
344                SqlGraphRelationship::RequiredByArg => format!(r#"color = "black""#),
345                SqlGraphRelationship::RequiredByReturn => {
346                    format!(r#"dir = "back", color = "black""#)
347                }
348            },
349            &|_graph, (_index, node)| {
350                match node {
351                    // Colors derived from https://www.schemecolor.com/touch-of-creativity.php
352                    SqlGraphEntity::Schema(_item) => format!(
353                        "label = \"{}\", weight = 6, shape = \"tab\"",
354                        node.dot_identifier()
355                    ),
356                    SqlGraphEntity::Function(_item) => format!(
357                        "label = \"{}\", penwidth = 0, style = \"filled\", fillcolor = \"#ADC7C6\", weight = 4, shape = \"box\"",
358                        node.dot_identifier()
359                    ),
360                    SqlGraphEntity::Type(_item) => format!(
361                        "label = \"{}\", penwidth = 0, style = \"filled\", fillcolor = \"#AE9BBD\", weight = 5, shape = \"oval\"",
362                        node.dot_identifier()
363                    ),
364                    SqlGraphEntity::BuiltinType(_item) => format!(
365                        "label = \"{}\", shape = \"plain\"",
366                        node.dot_identifier()
367                    ),
368                    SqlGraphEntity::Enum(_item) => format!(
369                        "label = \"{}\", penwidth = 0, style = \"filled\", fillcolor = \"#C9A7C8\", weight = 5, shape = \"oval\"",
370                        node.dot_identifier()
371                    ),
372                    SqlGraphEntity::Ord(_item) => format!(
373                        "label = \"{}\", penwidth = 0, style = \"filled\", fillcolor = \"#FFCFD3\", weight = 5, shape = \"diamond\"",
374                        node.dot_identifier()
375                    ),
376                    SqlGraphEntity::Hash(_item) => format!(
377                        "label = \"{}\", penwidth = 0, style = \"filled\", fillcolor = \"#FFE4E0\", weight = 5, shape = \"diamond\"",
378                        node.dot_identifier()
379                    ),
380                    SqlGraphEntity::Aggregate(_item) => format!(
381                        "label = \"{}\", penwidth = 0, style = \"filled\", fillcolor = \"#FFE4E0\", weight = 5, shape = \"diamond\"",
382                        node.dot_identifier()
383                    ),
384                    SqlGraphEntity::Trigger(_item) => format!(
385                        "label = \"{}\", penwidth = 0, style = \"filled\", fillcolor = \"#FFE4E0\", weight = 5, shape = \"diamond\"",
386                        node.dot_identifier()
387                    ),
388                    SqlGraphEntity::CustomSql(_item) => format!(
389                        "label = \"{}\", weight = 3, shape = \"signature\"",
390                        node.dot_identifier()
391                    ),
392                    SqlGraphEntity::ExtensionRoot(_item) => format!(
393                        "label = \"{}\", shape = \"cylinder\"",
394                        node.dot_identifier()
395                    ),
396                }
397            },
398        );
399        let path = Path::new(file.as_ref());
400
401        let parent = path.parent();
402        if let Some(parent) = parent {
403            create_dir_all(parent)?;
404        }
405        let mut out = File::create(path)?;
406        write!(out, "{:?}", generated)?;
407        Ok(())
408    }
409
410    pub fn schema_alias_of(&self, item_index: &NodeIndex) -> Option<String> {
411        self.graph
412            .neighbors_undirected(*item_index)
413            .flat_map(|neighbor_index| match &self.graph[neighbor_index] {
414                SqlGraphEntity::Schema(s) => Some(String::from(s.name)),
415                SqlGraphEntity::ExtensionRoot(control) => {
416                    if !control.relocatable {
417                        control.schema.clone()
418                    } else {
419                        Some(String::from("@extname@"))
420                    }
421                }
422                _ => None,
423            })
424            .next()
425    }
426
427    pub fn schema_prefix_for(&self, target: &NodeIndex) -> String {
428        self.schema_alias_of(target)
429            .map(|v| (v + ".").to_string())
430            .unwrap_or_else(|| "".to_string())
431    }
432
433    #[instrument(level = "error", skip(self))]
434    pub fn to_sql(&self) -> eyre::Result<String> {
435        let mut full_sql = String::new();
436        for step_id in petgraph::algo::toposort(&self.graph, None).map_err(|e| {
437            eyre!("Failed to toposort SQL entities, node with cycle: {:?}", self.graph[e.node_id()])
438        })? {
439            let step = &self.graph[step_id];
440
441            let sql = step.to_sql(self)?;
442
443            if !sql.is_empty() {
444                full_sql.push_str(&sql);
445                full_sql.push('\n');
446            }
447        }
448        Ok(full_sql)
449    }
450
451    pub fn has_sql_declared_entity(&self, identifier: &SqlDeclared) -> Option<&SqlDeclaredEntity> {
452        self.extension_sqls.iter().find_map(|(item, _index)| {
453            let retval = item.creates.iter().find_map(|create_entity| {
454                if create_entity.has_sql_declared_entity(identifier) {
455                    Some(create_entity)
456                } else {
457                    None
458                }
459            });
460            retval
461        })
462    }
463
464    pub fn source_only_to_sql_type(&self, ty_source: &str) -> Option<String> {
465        self.source_mappings.get(ty_source).map(|f| f.sql.clone())
466    }
467
468    pub fn get_module_pathname(&self) -> String {
469        return if self.versioned_so {
470            let extname = &self.extension_name;
471            let extver = &self.control.default_version;
472            // Note: versioned so-name format must agree with cargo pgx
473            format!("$libdir/{}-{}", extname, extver)
474        } else {
475            String::from("MODULE_PATHNAME")
476        };
477    }
478}
479
480#[tracing::instrument(level = "error", skip_all)]
481fn build_base_edges(
482    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
483    index: NodeIndex,
484    root: NodeIndex,
485    bootstrap: Option<NodeIndex>,
486    finalize: Option<NodeIndex>,
487) {
488    graph.add_edge(root, index, SqlGraphRelationship::RequiredBy);
489    if let Some(bootstrap) = bootstrap {
490        graph.add_edge(bootstrap, index, SqlGraphRelationship::RequiredBy);
491    }
492    if let Some(finalize) = finalize {
493        graph.add_edge(index, finalize, SqlGraphRelationship::RequiredBy);
494    }
495}
496
497#[tracing::instrument(level = "error", skip_all)]
498fn initialize_extension_sqls<'a>(
499    graph: &'a mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
500    root: NodeIndex,
501    extension_sqls: Vec<ExtensionSqlEntity>,
502) -> eyre::Result<(HashMap<ExtensionSqlEntity, NodeIndex>, Option<NodeIndex>, Option<NodeIndex>)> {
503    let mut bootstrap = None;
504    let mut finalize = None;
505    let mut mapped_extension_sqls = HashMap::default();
506    for item in extension_sqls {
507        let entity: SqlGraphEntity = item.clone().into();
508        let index = graph.add_node(entity);
509        mapped_extension_sqls.insert(item.clone(), index);
510
511        if item.bootstrap {
512            if let Some(exiting_index) = bootstrap {
513                let existing: &SqlGraphEntity = &graph[exiting_index];
514                return Err(eyre!(
515                    "Cannot have multiple `extension_sql!()` with `bootstrap` positioning, found `{}`, other was `{}`",
516                    item.rust_identifier(),
517                    existing.rust_identifier(),
518                ));
519            }
520            bootstrap = Some(index)
521        }
522        if item.finalize {
523            if let Some(exiting_index) = finalize {
524                let existing: &SqlGraphEntity = &graph[exiting_index];
525                return Err(eyre!(
526                    "Cannot have multiple `extension_sql!()` with `finalize` positioning, found `{}`, other was `{}`",
527                    item.rust_identifier(),
528                    existing.rust_identifier(),
529                ));
530            }
531            finalize = Some(index)
532        }
533    }
534    for (item, index) in &mapped_extension_sqls {
535        graph.add_edge(root, *index, SqlGraphRelationship::RequiredBy);
536        if !item.bootstrap {
537            if let Some(bootstrap) = bootstrap {
538                graph.add_edge(bootstrap, *index, SqlGraphRelationship::RequiredBy);
539            }
540        }
541        if !item.finalize {
542            if let Some(finalize) = finalize {
543                graph.add_edge(*index, finalize, SqlGraphRelationship::RequiredBy);
544            }
545        }
546    }
547    Ok((mapped_extension_sqls, bootstrap, finalize))
548}
549
550#[tracing::instrument(level = "error", skip_all)]
551/// A best effort attempt to find the related [`NodeIndex`] for some [`PositioningRef`].
552pub fn find_positioning_ref_target<'a>(
553    positioning_ref: &'a PositioningRef,
554    types: &'a HashMap<PostgresTypeEntity, NodeIndex>,
555    enums: &'a HashMap<PostgresEnumEntity, NodeIndex>,
556    externs: &'a HashMap<PgExternEntity, NodeIndex>,
557    schemas: &'a HashMap<SchemaEntity, NodeIndex>,
558    extension_sqls: &'a HashMap<ExtensionSqlEntity, NodeIndex>,
559    triggers: &'a HashMap<PgTriggerEntity, NodeIndex>,
560) -> Option<&'a NodeIndex> {
561    match positioning_ref {
562        PositioningRef::FullPath(path) => {
563            // The best we can do here is a fuzzy search.
564            let segments = path.split("::").collect::<Vec<_>>();
565            let last_segment = segments.last().expect("Expected at least one segment.");
566            let rest = &segments[..segments.len() - 1];
567            let module_path = rest.join("::");
568
569            for (other, other_index) in types {
570                if *last_segment == other.name && other.module_path.ends_with(&module_path) {
571                    return Some(&other_index);
572                }
573            }
574            for (other, other_index) in enums {
575                if last_segment == &other.name && other.module_path.ends_with(&module_path) {
576                    return Some(&other_index);
577                }
578            }
579            for (other, other_index) in externs {
580                if *last_segment == other.unaliased_name
581                    && other.module_path.ends_with(&module_path)
582                {
583                    return Some(&other_index);
584                }
585            }
586            for (other, other_index) in schemas {
587                if other.module_path.ends_with(path) {
588                    return Some(&other_index);
589                }
590            }
591
592            for (other, other_index) in triggers {
593                if last_segment == &other.function_name && other.module_path.ends_with(&module_path)
594                {
595                    return Some(&other_index);
596                }
597            }
598        }
599        PositioningRef::Name(name) => {
600            for (other, other_index) in extension_sqls {
601                if other.name == *name {
602                    return Some(&other_index);
603                }
604            }
605        }
606    };
607    None
608}
609
610#[tracing::instrument(level = "error", skip_all)]
611fn connect_extension_sqls(
612    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
613    extension_sqls: &HashMap<ExtensionSqlEntity, NodeIndex>,
614    schemas: &HashMap<SchemaEntity, NodeIndex>,
615    types: &HashMap<PostgresTypeEntity, NodeIndex>,
616    enums: &HashMap<PostgresEnumEntity, NodeIndex>,
617    externs: &HashMap<PgExternEntity, NodeIndex>,
618    triggers: &HashMap<PgTriggerEntity, NodeIndex>,
619) -> eyre::Result<()> {
620    for (item, &index) in extension_sqls {
621        make_schema_connection(
622            graph,
623            "Extension SQL",
624            index,
625            &item.rust_identifier(),
626            item.module_path,
627            schemas,
628        );
629
630        for requires in &item.requires {
631            if let Some(target) = find_positioning_ref_target(
632                requires,
633                types,
634                enums,
635                externs,
636                schemas,
637                extension_sqls,
638                triggers,
639            ) {
640                tracing::debug!(from = %item.rust_identifier(), to = ?graph[*target].rust_identifier(), "Adding ExtensionSQL after positioning ref target");
641                graph.add_edge(*target, index, SqlGraphRelationship::RequiredBy);
642            } else {
643                return Err(eyre!(
644                    "Could not find `requires` target of `{}`{}: {}",
645                    item.rust_identifier(),
646                    if let (Some(file), Some(line)) = (item.file(), item.line()) {
647                        format!(" ({}:{})", file, line)
648                    } else {
649                        "".to_string()
650                    },
651                    match requires {
652                        PositioningRef::FullPath(path) => path.to_string(),
653                        PositioningRef::Name(name) => format!(r#""{}""#, name),
654                    },
655                ));
656            }
657        }
658    }
659    Ok(())
660}
661
662#[tracing::instrument(level = "error", skip_all)]
663fn initialize_schemas(
664    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
665    bootstrap: Option<NodeIndex>,
666    finalize: Option<NodeIndex>,
667    schemas: Vec<SchemaEntity>,
668) -> eyre::Result<HashMap<SchemaEntity, NodeIndex>> {
669    let mut mapped_schemas = HashMap::default();
670    for item in schemas {
671        let entity = item.clone().into();
672        let index = graph.add_node(entity);
673        mapped_schemas.insert(item, index);
674        if let Some(bootstrap) = bootstrap {
675            graph.add_edge(bootstrap, index, SqlGraphRelationship::RequiredBy);
676        }
677        if let Some(finalize) = finalize {
678            graph.add_edge(index, finalize, SqlGraphRelationship::RequiredBy);
679        }
680    }
681    Ok(mapped_schemas)
682}
683
684#[tracing::instrument(level = "error", skip_all)]
685fn connect_schemas(
686    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
687    schemas: &HashMap<SchemaEntity, NodeIndex>,
688    root: NodeIndex,
689) {
690    for (_item, &index) in schemas {
691        graph.add_edge(root, index, SqlGraphRelationship::RequiredBy);
692    }
693}
694
695#[tracing::instrument(level = "error", skip_all)]
696fn initialize_enums(
697    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
698    root: NodeIndex,
699    bootstrap: Option<NodeIndex>,
700    finalize: Option<NodeIndex>,
701    enums: Vec<PostgresEnumEntity>,
702) -> eyre::Result<HashMap<PostgresEnumEntity, NodeIndex>> {
703    let mut mapped_enums = HashMap::default();
704    for item in enums {
705        let entity: SqlGraphEntity = item.clone().into();
706        let index = graph.add_node(entity);
707        mapped_enums.insert(item, index);
708        build_base_edges(graph, index, root, bootstrap, finalize);
709    }
710    Ok(mapped_enums)
711}
712
713#[tracing::instrument(level = "error", skip_all)]
714fn connect_enums(
715    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
716    enums: &HashMap<PostgresEnumEntity, NodeIndex>,
717    schemas: &HashMap<SchemaEntity, NodeIndex>,
718) {
719    for (item, &index) in enums {
720        make_schema_connection(
721            graph,
722            "Enum",
723            index,
724            &item.rust_identifier(),
725            item.module_path,
726            schemas,
727        );
728    }
729}
730
731#[tracing::instrument(level = "error", skip_all)]
732fn initialize_types(
733    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
734    root: NodeIndex,
735    bootstrap: Option<NodeIndex>,
736    finalize: Option<NodeIndex>,
737    types: Vec<PostgresTypeEntity>,
738) -> eyre::Result<HashMap<PostgresTypeEntity, NodeIndex>> {
739    let mut mapped_types = HashMap::default();
740    for item in types {
741        let entity = item.clone().into();
742        let index = graph.add_node(entity);
743        mapped_types.insert(item, index);
744        build_base_edges(graph, index, root, bootstrap, finalize);
745    }
746    Ok(mapped_types)
747}
748
749#[tracing::instrument(level = "error", skip_all)]
750fn connect_types(
751    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
752    types: &HashMap<PostgresTypeEntity, NodeIndex>,
753    schemas: &HashMap<SchemaEntity, NodeIndex>,
754) {
755    for (item, &index) in types {
756        make_schema_connection(
757            graph,
758            "Type",
759            index,
760            &item.rust_identifier(),
761            item.module_path,
762            schemas,
763        );
764    }
765}
766
767#[tracing::instrument(level = "error", skip_all)]
768fn initialize_externs(
769    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
770    root: NodeIndex,
771    bootstrap: Option<NodeIndex>,
772    finalize: Option<NodeIndex>,
773    externs: Vec<PgExternEntity>,
774    mapped_types: &HashMap<PostgresTypeEntity, NodeIndex>,
775    mapped_enums: &HashMap<PostgresEnumEntity, NodeIndex>,
776) -> eyre::Result<(HashMap<PgExternEntity, NodeIndex>, HashMap<String, NodeIndex>)> {
777    let mut mapped_externs = HashMap::default();
778    let mut mapped_builtin_types = HashMap::default();
779    for item in externs {
780        let entity: SqlGraphEntity = item.clone().into();
781        let index = graph.add_node(entity.clone());
782        mapped_externs.insert(item.clone(), index);
783        build_base_edges(graph, index, root, bootstrap, finalize);
784
785        for arg in &item.fn_args {
786            let mut found = false;
787            for (ty_item, &_ty_index) in mapped_types {
788                if ty_item.id_matches(&arg.used_ty.ty_id) {
789                    found = true;
790                    break;
791                }
792            }
793            for (ty_item, &_ty_index) in mapped_enums {
794                if ty_item.id_matches(&arg.used_ty.ty_id) {
795                    found = true;
796                    break;
797                }
798            }
799            if !found {
800                mapped_builtin_types.entry(arg.used_ty.full_path.to_string()).or_insert_with(
801                    || {
802                        graph.add_node(SqlGraphEntity::BuiltinType(
803                            arg.used_ty.full_path.to_string(),
804                        ))
805                    },
806                );
807            }
808        }
809
810        match &item.fn_return {
811            PgExternReturnEntity::None | PgExternReturnEntity::Trigger => (),
812            PgExternReturnEntity::Type { ty, .. } | PgExternReturnEntity::SetOf { ty, .. } => {
813                let mut found = false;
814                for (ty_item, &_ty_index) in mapped_types {
815                    if ty_item.id_matches(&ty.ty_id) {
816                        found = true;
817                        break;
818                    }
819                }
820                for (ty_item, &_ty_index) in mapped_enums {
821                    if ty_item.id_matches(&ty.ty_id) {
822                        found = true;
823                        break;
824                    }
825                }
826                if !found {
827                    mapped_builtin_types.entry(ty.full_path.to_string()).or_insert_with(|| {
828                        graph.add_node(SqlGraphEntity::BuiltinType(ty.full_path.to_string()))
829                    });
830                }
831            }
832            PgExternReturnEntity::Iterated { tys: iterated_returns, optional: _ } => {
833                for PgExternReturnEntityIteratedItem { ty: return_ty_entity, .. } in
834                    iterated_returns
835                {
836                    let mut found = false;
837                    for (ty_item, &_ty_index) in mapped_types {
838                        if ty_item.id_matches(&return_ty_entity.ty_id) {
839                            found = true;
840                            break;
841                        }
842                    }
843                    for (ty_item, &_ty_index) in mapped_enums {
844                        if ty_item.id_matches(&return_ty_entity.ty_id) {
845                            found = true;
846                            break;
847                        }
848                    }
849                    if !found {
850                        mapped_builtin_types
851                            .entry(return_ty_entity.ty_source.to_string())
852                            .or_insert_with(|| {
853                                graph.add_node(SqlGraphEntity::BuiltinType(
854                                    return_ty_entity.ty_source.to_string(),
855                                ))
856                            });
857                    }
858                }
859            }
860        }
861    }
862    Ok((mapped_externs, mapped_builtin_types))
863}
864
865#[tracing::instrument(level = "error", skip_all)]
866fn connect_externs(
867    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
868    externs: &HashMap<PgExternEntity, NodeIndex>,
869    hashes: &HashMap<PostgresHashEntity, NodeIndex>,
870    schemas: &HashMap<SchemaEntity, NodeIndex>,
871    types: &HashMap<PostgresTypeEntity, NodeIndex>,
872    enums: &HashMap<PostgresEnumEntity, NodeIndex>,
873    builtin_types: &HashMap<String, NodeIndex>,
874    extension_sqls: &HashMap<ExtensionSqlEntity, NodeIndex>,
875    triggers: &HashMap<PgTriggerEntity, NodeIndex>,
876) -> eyre::Result<()> {
877    for (item, &index) in externs {
878        let mut found_schema_declaration = false;
879        for extern_attr in &item.extern_attrs {
880            match extern_attr {
881                crate::ExternArgs::Requires(requirements) => {
882                    for requires in requirements {
883                        if let Some(target) = find_positioning_ref_target(
884                            requires,
885                            types,
886                            enums,
887                            externs,
888                            schemas,
889                            extension_sqls,
890                            triggers,
891                        ) {
892                            tracing::debug!(from = %item.rust_identifier(), to = %graph[*target].rust_identifier(), "Adding Extern after positioning ref target");
893                            graph.add_edge(*target, index, SqlGraphRelationship::RequiredBy);
894                        } else {
895                            return Err(eyre!("Could not find `requires` target: {:?}", requires));
896                        }
897                    }
898                }
899                crate::ExternArgs::Schema(declared_schema_name) => {
900                    for (schema, schema_index) in schemas {
901                        if schema.name == declared_schema_name {
902                            tracing::debug!(from = ?item.rust_identifier(), to = schema.module_path, "Adding Extern after Schema edge.");
903                            graph.add_edge(*schema_index, index, SqlGraphRelationship::RequiredBy);
904                            found_schema_declaration = true;
905                        }
906                    }
907                    if !found_schema_declaration {
908                        return Err(eyre!("Got manual `schema = \"{declared_schema_name}\"` setting, but that schema did not exist."));
909                    }
910                }
911                _ => (),
912            }
913        }
914
915        if !found_schema_declaration {
916            make_schema_connection(
917                graph,
918                "Extern",
919                index,
920                &item.rust_identifier(),
921                item.module_path,
922                schemas,
923            );
924        }
925
926        // The hash function must be defined after the {typename}_eq function.
927        for (hash_item, &hash_index) in hashes {
928            if item.module_path == hash_item.module_path
929                && item.name == hash_item.name.to_lowercase() + "_eq"
930            {
931                tracing::debug!(from = hash_item.full_path, to = ?item.full_path, "Adding Hash after Extern edge");
932                graph.add_edge(index, hash_index, SqlGraphRelationship::RequiredBy);
933            }
934        }
935
936        for arg in &item.fn_args {
937            let mut found = false;
938
939            for (ty_item, &ty_index) in types {
940                if ty_item.id_matches(&arg.used_ty.ty_id) {
941                    tracing::debug!(from = %item.rust_identifier(), to = %ty_item.rust_identifier(), "Adding Extern after Type (due to argument) edge");
942                    graph.add_edge(ty_index, index, SqlGraphRelationship::RequiredByArg);
943                    found = true;
944                    break;
945                }
946            }
947            if !found {
948                for (enum_item, &enum_index) in enums {
949                    if enum_item.id_matches(&arg.used_ty.ty_id) {
950                        tracing::debug!(from = %item.rust_identifier(), to = %enum_item.rust_identifier(), "Adding Extern after Enum (due to argument) edge");
951                        graph.add_edge(enum_index, index, SqlGraphRelationship::RequiredByArg);
952                        found = true;
953                        break;
954                    }
955                }
956            }
957            if !found {
958                let builtin_index = builtin_types
959                    .get(arg.used_ty.full_path)
960                    .expect(&format!("Could not fetch Builtin Type {}.", arg.used_ty.full_path));
961                tracing::debug!(from = %item.rust_identifier(), to = %arg.rust_identifier(), "Adding Extern(arg) after BuiltIn Type (due to argument) edge");
962                graph.add_edge(*builtin_index, index, SqlGraphRelationship::RequiredByArg);
963            }
964            if !found {
965                for (ext_item, ext_index) in extension_sqls {
966                    if let Some(_) = ext_item.has_sql_declared_entity(&SqlDeclared::Type(
967                        arg.used_ty.full_path.to_string(),
968                    )) {
969                        tracing::debug!(from = %item.rust_identifier(), to = %arg.rust_identifier(), "Adding Extern(arg) after Extension SQL (due to argument) edge");
970                        graph.add_edge(*ext_index, index, SqlGraphRelationship::RequiredByArg);
971                    } else if let Some(_) = ext_item.has_sql_declared_entity(&SqlDeclared::Enum(
972                        arg.used_ty.full_path.to_string(),
973                    )) {
974                        tracing::debug!(from = %item.rust_identifier(), to = %arg.rust_identifier(), "Adding Extern(arg) after Extension SQL (due to argument) edge");
975                        graph.add_edge(*ext_index, index, SqlGraphRelationship::RequiredByArg);
976                    }
977                }
978            }
979        }
980        match &item.fn_return {
981            PgExternReturnEntity::None | PgExternReturnEntity::Trigger => (),
982            PgExternReturnEntity::Type { ty, .. } | PgExternReturnEntity::SetOf { ty, .. } => {
983                let mut found = false;
984                for (ty_item, &ty_index) in types {
985                    if ty_item.id_matches(&ty.ty_id) {
986                        tracing::debug!(from = %item.rust_identifier(), to = %ty_item.rust_identifier(), "Adding Extern after Type (due to return) edge");
987                        graph.add_edge(ty_index, index, SqlGraphRelationship::RequiredByReturn);
988                        found = true;
989                        break;
990                    }
991                }
992                if !found {
993                    for (ty_item, &ty_index) in enums {
994                        if ty_item.id_matches(&ty.ty_id) {
995                            tracing::debug!(from = %item.rust_identifier(), to = %ty_item.rust_identifier(), "Adding Extern after Enum (due to return) edge");
996                            graph.add_edge(ty_index, index, SqlGraphRelationship::RequiredByReturn);
997                            found = true;
998                            break;
999                        }
1000                    }
1001                }
1002                if !found {
1003                    let builtin_index = builtin_types
1004                        .get(&ty.full_path.to_string())
1005                        .expect(&format!("Could not fetch Builtin Type {}.", ty.full_path));
1006                    tracing::debug!(from = ?item.full_path, to = %ty.full_path, "Adding Extern(return) after BuiltIn Type (due to return) edge");
1007                    graph.add_edge(*builtin_index, index, SqlGraphRelationship::RequiredByReturn);
1008                }
1009                if !found {
1010                    for (ext_item, ext_index) in extension_sqls {
1011                        if let Some(_) = ext_item
1012                            .has_sql_declared_entity(&SqlDeclared::Type(ty.full_path.to_string()))
1013                        {
1014                            tracing::debug!(from = %item.rust_identifier(), to = ty.full_path, "Adding Extern(arg) after Extension SQL (due to argument) edge");
1015                            graph.add_edge(*ext_index, index, SqlGraphRelationship::RequiredByArg);
1016                        } else if let Some(_) = ext_item
1017                            .has_sql_declared_entity(&SqlDeclared::Enum(ty.full_path.to_string()))
1018                        {
1019                            tracing::debug!(from = %item.rust_identifier(), to = ty.full_path, "Adding Extern(arg) after Extension SQL (due to argument) edge");
1020                            graph.add_edge(*ext_index, index, SqlGraphRelationship::RequiredByArg);
1021                        }
1022                    }
1023                }
1024            }
1025            PgExternReturnEntity::Iterated { tys: iterated_returns, optional: _ } => {
1026                for PgExternReturnEntityIteratedItem { ty: type_entity, .. } in iterated_returns {
1027                    let mut found = false;
1028                    for (ty_item, &ty_index) in types {
1029                        if ty_item.id_matches(&type_entity.ty_id) {
1030                            tracing::debug!(from = %item.rust_identifier(), to = %ty_item.rust_identifier(), "Adding Extern after Type (due to return) edge");
1031                            graph.add_edge(ty_index, index, SqlGraphRelationship::RequiredByReturn);
1032                            found = true;
1033                            break;
1034                        }
1035                    }
1036                    if !found {
1037                        for (ty_item, &ty_index) in enums {
1038                            if ty_item.id_matches(&type_entity.ty_id) {
1039                                tracing::debug!(from = %item.rust_identifier(), to = %ty_item.rust_identifier(), "Adding Extern after Enum (due to return) edge");
1040                                graph.add_edge(
1041                                    ty_index,
1042                                    index,
1043                                    SqlGraphRelationship::RequiredByReturn,
1044                                );
1045                                found = true;
1046                                break;
1047                            }
1048                        }
1049                    }
1050                    if !found {
1051                        let builtin_index =
1052                            builtin_types.get(&type_entity.ty_source.to_string()).expect(&format!(
1053                                "Could not fetch Builtin Type {}.",
1054                                type_entity.ty_source,
1055                            ));
1056                        tracing::debug!(from = %item.rust_identifier(), to = type_entity.ty_source, "Adding Extern after BuiltIn Type (due to return) edge");
1057                        graph.add_edge(
1058                            *builtin_index,
1059                            index,
1060                            SqlGraphRelationship::RequiredByReturn,
1061                        );
1062                    }
1063                    if !found {
1064                        for (ext_item, ext_index) in extension_sqls {
1065                            if let Some(_) = ext_item.has_sql_declared_entity(&SqlDeclared::Type(
1066                                type_entity.ty_source.to_string(),
1067                            )) {
1068                                tracing::debug!(from = %item.rust_identifier(), to = %ext_item.rust_identifier(), "Adding Extern(arg) after Extension SQL (due to argument) edge");
1069                                graph.add_edge(
1070                                    *ext_index,
1071                                    index,
1072                                    SqlGraphRelationship::RequiredByArg,
1073                                );
1074                            } else if let Some(_) = ext_item.has_sql_declared_entity(
1075                                &SqlDeclared::Enum(type_entity.ty_source.to_string()),
1076                            ) {
1077                                tracing::debug!(from = %item.rust_identifier(), to = %ext_item.rust_identifier(), "Adding Extern(arg) after Extension SQL (due to argument) edge");
1078                                graph.add_edge(
1079                                    *ext_index,
1080                                    index,
1081                                    SqlGraphRelationship::RequiredByArg,
1082                                );
1083                            }
1084                        }
1085                    }
1086                }
1087            }
1088        }
1089    }
1090    Ok(())
1091}
1092
1093#[tracing::instrument(level = "info", skip_all)]
1094fn initialize_ords(
1095    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
1096    root: NodeIndex,
1097    bootstrap: Option<NodeIndex>,
1098    finalize: Option<NodeIndex>,
1099    ords: Vec<PostgresOrdEntity>,
1100) -> eyre::Result<HashMap<PostgresOrdEntity, NodeIndex>> {
1101    let mut mapped_ords = HashMap::default();
1102    for item in ords {
1103        let entity = item.clone().into();
1104        let index = graph.add_node(entity);
1105        mapped_ords.insert(item.clone(), index);
1106        build_base_edges(graph, index, root, bootstrap, finalize);
1107    }
1108    Ok(mapped_ords)
1109}
1110
1111#[tracing::instrument(level = "info", skip_all)]
1112fn connect_ords(
1113    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
1114    ords: &HashMap<PostgresOrdEntity, NodeIndex>,
1115    schemas: &HashMap<SchemaEntity, NodeIndex>,
1116    types: &HashMap<PostgresTypeEntity, NodeIndex>,
1117    enums: &HashMap<PostgresEnumEntity, NodeIndex>,
1118    externs: &HashMap<PgExternEntity, NodeIndex>,
1119) {
1120    for (item, &index) in ords {
1121        make_schema_connection(
1122            graph,
1123            "Ord",
1124            index,
1125            &item.rust_identifier(),
1126            item.module_path,
1127            schemas,
1128        );
1129
1130        make_type_or_enum_connection(
1131            graph,
1132            "Ord",
1133            index,
1134            &item.rust_identifier(),
1135            &item.id,
1136            types,
1137            enums,
1138        );
1139
1140        // Make PostgresOrdEntities (which will be translated into `CREATE OPERATOR CLASS` statements) depend
1141        // on the operators which they will reference. For example, a pgx-defined Postgres type `parakeet`
1142        // which has `#[derive(PostgresOrd)]` will emit a `parakeet_btree_ops` operator class, which references
1143        // a definition of a < operator (among others) on the `parakeet` type. This code should ensure that the
1144        // < operator (along with all the others) is emitted before the `OPERATOR CLASS` itself.
1145
1146        for (extern_item, &extern_index) in externs {
1147            let fn_matches = |fn_name| {
1148                item.module_path == extern_item.module_path && extern_item.name == fn_name
1149            };
1150            let cmp_fn_matches = fn_matches(item.cmp_fn_name());
1151            let lt_fn_matches = fn_matches(item.lt_fn_name());
1152            let lte_fn_matches = fn_matches(item.le_fn_name());
1153            let eq_fn_matches = fn_matches(item.eq_fn_name());
1154            let gt_fn_matches = fn_matches(item.gt_fn_name());
1155            let gte_fn_matches = fn_matches(item.ge_fn_name());
1156            if cmp_fn_matches
1157                || lt_fn_matches
1158                || lte_fn_matches
1159                || eq_fn_matches
1160                || gt_fn_matches
1161                || gte_fn_matches
1162            {
1163                tracing::debug!(from = ?item.full_path, to = extern_item.full_path, "Adding Ord after Extern edge");
1164                graph.add_edge(extern_index, index, SqlGraphRelationship::RequiredBy);
1165            }
1166        }
1167    }
1168}
1169
1170#[tracing::instrument(level = "info", skip_all)]
1171fn initialize_hashes(
1172    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
1173    root: NodeIndex,
1174    bootstrap: Option<NodeIndex>,
1175    finalize: Option<NodeIndex>,
1176    hashes: Vec<PostgresHashEntity>,
1177) -> eyre::Result<HashMap<PostgresHashEntity, NodeIndex>> {
1178    let mut mapped_hashes = HashMap::default();
1179    for item in hashes {
1180        let entity: SqlGraphEntity = item.clone().into();
1181        let index = graph.add_node(entity);
1182        mapped_hashes.insert(item, index);
1183        build_base_edges(graph, index, root, bootstrap, finalize);
1184    }
1185    Ok(mapped_hashes)
1186}
1187
1188#[tracing::instrument(level = "info", skip_all)]
1189fn connect_hashes(
1190    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
1191    hashes: &HashMap<PostgresHashEntity, NodeIndex>,
1192    schemas: &HashMap<SchemaEntity, NodeIndex>,
1193    types: &HashMap<PostgresTypeEntity, NodeIndex>,
1194    enums: &HashMap<PostgresEnumEntity, NodeIndex>,
1195    externs: &HashMap<PgExternEntity, NodeIndex>,
1196) {
1197    for (item, &index) in hashes {
1198        make_schema_connection(
1199            graph,
1200            "Hash",
1201            index,
1202            &item.rust_identifier(),
1203            item.module_path,
1204            schemas,
1205        );
1206
1207        make_type_or_enum_connection(
1208            graph,
1209            "Hash",
1210            index,
1211            &item.rust_identifier(),
1212            &item.id,
1213            types,
1214            enums,
1215        );
1216
1217        for (extern_item, &extern_index) in externs {
1218            let hash_fn_name = item.fn_name();
1219            let hash_fn_matches =
1220                item.module_path == extern_item.module_path && extern_item.name == hash_fn_name;
1221
1222            if hash_fn_matches {
1223                tracing::debug!(from = ?item.full_path, to = extern_item.full_path, "Adding Hash after Extern edge");
1224                graph.add_edge(extern_index, index, SqlGraphRelationship::RequiredBy);
1225                break;
1226            }
1227        }
1228    }
1229}
1230
1231#[tracing::instrument(level = "info", skip_all)]
1232fn initialize_aggregates(
1233    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
1234    root: NodeIndex,
1235    bootstrap: Option<NodeIndex>,
1236    finalize: Option<NodeIndex>,
1237    aggregates: Vec<PgAggregateEntity>,
1238    mapped_builtin_types: &mut HashMap<String, NodeIndex>,
1239    mapped_enums: &HashMap<PostgresEnumEntity, NodeIndex>,
1240    mapped_types: &HashMap<PostgresTypeEntity, NodeIndex>,
1241) -> eyre::Result<HashMap<PgAggregateEntity, NodeIndex>> {
1242    let mut mapped_aggregates = HashMap::default();
1243    for item in aggregates {
1244        let entity: SqlGraphEntity = item.clone().into();
1245        let index = graph.add_node(entity);
1246
1247        for arg in &item.args {
1248            let mut found = false;
1249            for (ty_item, &_ty_index) in mapped_types {
1250                if ty_item.id_matches(&arg.used_ty.ty_id) {
1251                    found = true;
1252                    break;
1253                }
1254            }
1255            for (ty_item, &_ty_index) in mapped_enums {
1256                if ty_item.id_matches(&arg.used_ty.ty_id) {
1257                    found = true;
1258                    break;
1259                }
1260            }
1261            if !found {
1262                mapped_builtin_types.entry(arg.used_ty.full_path.to_string()).or_insert_with(
1263                    || {
1264                        graph.add_node(SqlGraphEntity::BuiltinType(
1265                            arg.used_ty.full_path.to_string(),
1266                        ))
1267                    },
1268                );
1269            }
1270        }
1271
1272        mapped_aggregates.insert(item, index);
1273        build_base_edges(graph, index, root, bootstrap, finalize);
1274    }
1275    Ok(mapped_aggregates)
1276}
1277
1278#[tracing::instrument(level = "error", skip_all, fields(rust_identifier = %item.rust_identifier()))]
1279fn connect_aggregate(
1280    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
1281    item: &PgAggregateEntity,
1282    index: NodeIndex,
1283    schemas: &HashMap<SchemaEntity, NodeIndex>,
1284    types: &HashMap<PostgresTypeEntity, NodeIndex>,
1285    enums: &HashMap<PostgresEnumEntity, NodeIndex>,
1286    builtin_types: &HashMap<String, NodeIndex>,
1287    externs: &HashMap<PgExternEntity, NodeIndex>,
1288) -> eyre::Result<()> {
1289    make_schema_connection(
1290        graph,
1291        "Aggregate",
1292        index,
1293        &item.rust_identifier(),
1294        item.module_path,
1295        schemas,
1296    );
1297
1298    make_type_or_enum_connection(
1299        graph,
1300        "Aggregate",
1301        index,
1302        &item.rust_identifier(),
1303        &item.ty_id,
1304        types,
1305        enums,
1306    );
1307
1308    for arg in &item.args {
1309        let found = make_type_or_enum_connection(
1310            graph,
1311            "Aggregate",
1312            index,
1313            &item.rust_identifier(),
1314            &arg.used_ty.ty_id,
1315            types,
1316            enums,
1317        );
1318        if !found {
1319            let builtin_index = builtin_types
1320                .get(arg.used_ty.full_path)
1321                .expect(&format!("Could not fetch Builtin Type {}.", arg.used_ty.full_path));
1322            tracing::debug!(from = %item.rust_identifier(), to = %arg.used_ty.full_path, "Adding Aggregate after BuiltIn Type edge");
1323            graph.add_edge(*builtin_index, index, SqlGraphRelationship::RequiredByArg);
1324        }
1325    }
1326
1327    for arg in item.direct_args.as_ref().unwrap_or(&vec![]) {
1328        let found = make_type_or_enum_connection(
1329            graph,
1330            "Aggregate",
1331            index,
1332            &item.rust_identifier(),
1333            &arg.used_ty.ty_id,
1334            types,
1335            enums,
1336        );
1337        if !found {
1338            let builtin_index = builtin_types
1339                .get(arg.used_ty.full_path)
1340                .expect(&format!("Could not fetch Builtin Type {}.", arg.used_ty.full_path));
1341            tracing::debug!(from = %item.rust_identifier(), to = %arg.used_ty.full_path, "Adding Aggregate after BuiltIn Type edge");
1342            graph.add_edge(*builtin_index, index, SqlGraphRelationship::RequiredByArg);
1343        }
1344    }
1345
1346    if let Some(arg) = &item.mstype {
1347        let found = make_type_or_enum_connection(
1348            graph,
1349            "Aggregate",
1350            index,
1351            &item.rust_identifier(),
1352            &arg.ty_id,
1353            types,
1354            enums,
1355        );
1356        if !found {
1357            let builtin_index = builtin_types
1358                .get(arg.full_path)
1359                .expect(&format!("Could not fetch Builtin Type {}.", arg.full_path));
1360            tracing::debug!(from = %item.rust_identifier(), to = %arg.full_path, "Adding Aggregate after BuiltIn Type edge");
1361            graph.add_edge(*builtin_index, index, SqlGraphRelationship::RequiredByArg);
1362        }
1363    }
1364
1365    make_extern_connection(
1366        graph,
1367        "Aggregate",
1368        index,
1369        &item.rust_identifier(),
1370        &(item.module_path.to_string() + "::" + item.sfunc),
1371        externs,
1372    )?;
1373
1374    if let Some(value) = item.finalfunc {
1375        make_extern_connection(
1376            graph,
1377            "Aggregate",
1378            index,
1379            &item.rust_identifier(),
1380            &(item.module_path.to_string() + "::" + value),
1381            externs,
1382        )?;
1383    }
1384    if let Some(value) = item.combinefunc {
1385        make_extern_connection(
1386            graph,
1387            "Aggregate",
1388            index,
1389            &item.rust_identifier(),
1390            &(item.module_path.to_string() + "::" + value),
1391            externs,
1392        )?;
1393    }
1394    if let Some(value) = item.serialfunc {
1395        make_extern_connection(
1396            graph,
1397            "Aggregate",
1398            index,
1399            &item.rust_identifier(),
1400            &(item.module_path.to_string() + "::" + value),
1401            externs,
1402        )?;
1403    }
1404    if let Some(value) = item.deserialfunc {
1405        make_extern_connection(
1406            graph,
1407            "Aggregate",
1408            index,
1409            &item.rust_identifier(),
1410            &(item.module_path.to_string() + "::" + value),
1411            externs,
1412        )?;
1413    }
1414    if let Some(value) = item.msfunc {
1415        make_extern_connection(
1416            graph,
1417            "Aggregate",
1418            index,
1419            &item.rust_identifier(),
1420            &(item.module_path.to_string() + "::" + value),
1421            externs,
1422        )?;
1423    }
1424    if let Some(value) = item.minvfunc {
1425        make_extern_connection(
1426            graph,
1427            "Aggregate",
1428            index,
1429            &item.rust_identifier(),
1430            &(item.module_path.to_string() + "::" + value),
1431            externs,
1432        )?;
1433    }
1434    if let Some(value) = item.mfinalfunc {
1435        make_extern_connection(
1436            graph,
1437            "Aggregate",
1438            index,
1439            &item.rust_identifier(),
1440            &(item.module_path.to_string() + "::" + value),
1441            externs,
1442        )?;
1443    }
1444    if let Some(value) = item.sortop {
1445        make_extern_connection(
1446            graph,
1447            "Aggregate",
1448            index,
1449            &item.rust_identifier(),
1450            &(item.module_path.to_string() + "::" + value),
1451            externs,
1452        )?;
1453    }
1454    Ok(())
1455}
1456
1457#[tracing::instrument(level = "info", skip_all)]
1458fn connect_aggregates(
1459    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
1460    aggregates: &HashMap<PgAggregateEntity, NodeIndex>,
1461    schemas: &HashMap<SchemaEntity, NodeIndex>,
1462    types: &HashMap<PostgresTypeEntity, NodeIndex>,
1463    enums: &HashMap<PostgresEnumEntity, NodeIndex>,
1464    builtin_types: &HashMap<String, NodeIndex>,
1465    externs: &HashMap<PgExternEntity, NodeIndex>,
1466) -> eyre::Result<()> {
1467    for (item, &index) in aggregates {
1468        connect_aggregate(graph, item, index, schemas, types, enums, builtin_types, externs)?
1469    }
1470    Ok(())
1471}
1472
1473#[tracing::instrument(level = "info", skip_all)]
1474fn initialize_triggers(
1475    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
1476    root: NodeIndex,
1477    bootstrap: Option<NodeIndex>,
1478    finalize: Option<NodeIndex>,
1479    triggers: Vec<PgTriggerEntity>,
1480) -> eyre::Result<HashMap<PgTriggerEntity, NodeIndex>> {
1481    let mut mapped_triggers = HashMap::default();
1482    for item in triggers {
1483        let entity: SqlGraphEntity = item.clone().into();
1484        let index = graph.add_node(entity);
1485
1486        mapped_triggers.insert(item, index);
1487        build_base_edges(graph, index, root, bootstrap, finalize);
1488    }
1489    Ok(mapped_triggers)
1490}
1491
1492#[tracing::instrument(level = "info", skip_all)]
1493fn connect_triggers(
1494    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
1495    triggers: &HashMap<PgTriggerEntity, NodeIndex>,
1496    schemas: &HashMap<SchemaEntity, NodeIndex>,
1497) {
1498    for (item, &index) in triggers {
1499        make_schema_connection(
1500            graph,
1501            "Trigger",
1502            index,
1503            &item.rust_identifier(),
1504            item.module_path,
1505            schemas,
1506        );
1507    }
1508}
1509
1510#[tracing::instrument(level = "info", skip_all, fields(rust_identifier))]
1511fn make_schema_connection(
1512    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
1513    kind: &str,
1514    index: NodeIndex,
1515    rust_identifier: &str,
1516    module_path: &str,
1517    schemas: &HashMap<SchemaEntity, NodeIndex>,
1518) -> bool {
1519    let mut found = false;
1520    for (schema_item, &schema_index) in schemas {
1521        if module_path == schema_item.module_path {
1522            tracing::debug!(from = ?rust_identifier, to = schema_item.module_path, "Adding {kind} after Schema edge.", kind = kind);
1523            graph.add_edge(schema_index, index, SqlGraphRelationship::RequiredBy);
1524            found = true;
1525            break;
1526        }
1527    }
1528    found
1529}
1530
1531#[tracing::instrument(level = "error", skip_all, fields(%rust_identifier))]
1532fn make_extern_connection(
1533    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
1534    kind: &str,
1535    index: NodeIndex,
1536    rust_identifier: &str,
1537    full_path: &str,
1538    externs: &HashMap<PgExternEntity, NodeIndex>,
1539) -> eyre::Result<()> {
1540    let mut found = false;
1541    for (extern_item, &extern_index) in externs {
1542        if full_path == extern_item.full_path {
1543            tracing::debug!(from = ?rust_identifier, to = extern_item.full_path, "Adding {kind} after Extern edge.", kind = kind);
1544            graph.add_edge(extern_index, index, SqlGraphRelationship::RequiredBy);
1545            found = true;
1546            break;
1547        }
1548    }
1549    match found {
1550        true => Ok(()),
1551        false => Err(eyre!("Did not find connection `{full_path}` in {:#?}", {
1552            let mut paths = externs.iter().map(|(v, _)| v.full_path).collect::<Vec<_>>();
1553            paths.sort();
1554            paths
1555        })),
1556    }
1557}
1558
1559#[tracing::instrument(level = "info", skip_all, fields(rust_identifier))]
1560fn make_type_or_enum_connection(
1561    graph: &mut StableGraph<SqlGraphEntity, SqlGraphRelationship>,
1562    kind: &str,
1563    index: NodeIndex,
1564    rust_identifier: &str,
1565    ty_id: &TypeId,
1566    types: &HashMap<PostgresTypeEntity, NodeIndex>,
1567    enums: &HashMap<PostgresEnumEntity, NodeIndex>,
1568) -> bool {
1569    let mut found = false;
1570    for (ty_item, &ty_index) in types {
1571        if ty_item.id_matches(ty_id) {
1572            tracing::debug!(from = ?rust_identifier, to = ty_item.full_path, "Adding {kind} after Type edge.", kind = kind);
1573            graph.add_edge(ty_index, index, SqlGraphRelationship::RequiredBy);
1574            found = true;
1575            break;
1576        }
1577    }
1578    for (ty_item, &ty_index) in enums {
1579        if ty_item.id_matches(ty_id) {
1580            tracing::debug!(from = ?rust_identifier, to = ty_item.full_path, "Adding {kind} after Enum edge.", kind = kind);
1581            graph.add_edge(ty_index, index, SqlGraphRelationship::RequiredBy);
1582            found = true;
1583            break;
1584        }
1585    }
1586
1587    found
1588}
pgx_utils/sql_entity_graph/pgx_sql.rs

pgx_utils/sql_entity_graph/
pgx_sql.rs
