elefant_tools/storage/postgres/
postgres_instance_storage.rs

1use crate::postgres_client_wrapper::{FromPgChar, RowEnumExt};
2use crate::quoting::AllowedKeywordUsage;
3use crate::schema_reader::SchemaReader;
4use crate::storage::postgres::parallel_copy_destination::ParallelSafePostgresInstanceCopyDestinationStorage;
5use crate::storage::postgres::parallel_copy_source::ParallelSafePostgresInstanceCopySourceStorage;
6use crate::storage::postgres::sequential_copy_destination::SequentialSafePostgresInstanceCopyDestinationStorage;
7use crate::storage::postgres::sequential_copy_source::SequentialSafePostgresInstanceCopySourceStorage;
8use crate::{
9    BaseCopyTarget, CopyDestinationFactory, CopySourceFactory, DataFormat, ElefantToolsError,
10    IdentifierQuoter, PostgresClientWrapper, PostgresDatabase, SequentialOrParallel,
11    SupportedParallelism,
12};
13use elefant_client::PostgresDataRow;
14use std::collections::HashMap;
15use std::sync::Arc;
16use tracing::instrument;
17
18/// A CopyTarget for Postgres.
19pub struct PostgresInstanceStorage<'a> {
20    pub(crate) connection: &'a PostgresClientWrapper,
21    pub(crate) postgres_version: String,
22    pub(crate) identifier_quoter: Arc<IdentifierQuoter>,
23}
24
25impl<'a> PostgresInstanceStorage<'a> {
26    #[instrument(skip_all)]
27    pub async fn new(connection: &'a PostgresClientWrapper) -> crate::Result<Self> {
28        let postgres_version = connection.get_single_result("select version()").await?;
29
30        let keywords = connection
31            .get_results::<Keyword>(
32                "select word, catcode from pg_get_keywords() where catcode <> 'U'",
33            )
34            .await?;
35
36        let mut keyword_info = HashMap::new();
37
38        for keyword in keywords {
39            keyword_info.insert(
40                keyword.word,
41                AllowedKeywordUsage {
42                    column_name: keyword.category == KeywordType::AllowedInColumnName,
43                    type_or_function_name: keyword.category
44                        == KeywordType::AllowedInTypeOrFunctionName,
45                },
46            );
47        }
48
49        let quoter = IdentifierQuoter::new(keyword_info);
50
51        Ok(PostgresInstanceStorage {
52            connection,
53            postgres_version,
54            identifier_quoter: Arc::new(quoter),
55        })
56    }
57
58    pub fn get_identifier_quoter(&self) -> Arc<IdentifierQuoter> {
59        self.identifier_quoter.clone()
60    }
61}
62
63struct Keyword {
64    word: String,
65    category: KeywordType,
66}
67
68impl<'a> elefant_client::FromSqlRow<'a> for Keyword {
69    fn from_sql_row(
70        row: &'a PostgresDataRow<'_, '_>,
71    ) -> Result<Self, elefant_client::ElefantClientError> {
72        Ok(Keyword {
73            word: row.get(0)?,
74            category: row.try_get_enum_value(1)?,
75        })
76    }
77}
78
79#[derive(Eq, PartialEq, Debug)]
80enum KeywordType {
81    Unreserved,
82    AllowedInColumnName,
83    AllowedInTypeOrFunctionName,
84    Reserved,
85}
86
87impl FromPgChar for KeywordType {
88    fn from_pg_char(c: char) -> crate::Result<Self> {
89        match c {
90            'U' => Ok(KeywordType::Unreserved),
91            'C' => Ok(KeywordType::AllowedInColumnName),
92            'T' => Ok(KeywordType::AllowedInTypeOrFunctionName),
93            'R' => Ok(KeywordType::Reserved),
94            _ => Err(ElefantToolsError::InvalidKeywordType(c.to_string())),
95        }
96    }
97}
98
99impl BaseCopyTarget for PostgresInstanceStorage<'_> {
100    async fn supported_data_format(&self) -> crate::Result<Vec<DataFormat>> {
101        Ok(vec![
102            DataFormat::Text,
103            DataFormat::PostgresBinary {
104                postgres_version: Some(self.postgres_version.clone()),
105            },
106        ])
107    }
108}
109
110impl<'a> CopySourceFactory for PostgresInstanceStorage<'a> {
111    type SequentialSource = SequentialSafePostgresInstanceCopySourceStorage<'a>;
112    type ParallelSource = ParallelSafePostgresInstanceCopySourceStorage<'a>;
113
114    async fn get_introspection(&self) -> crate::Result<PostgresDatabase> {
115        let reader = SchemaReader::new(self.connection);
116        reader.introspect_database().await
117    }
118
119    async fn create_source(
120        &self,
121    ) -> crate::Result<SequentialOrParallel<Self::SequentialSource, Self::ParallelSource>> {
122        let parallel = ParallelSafePostgresInstanceCopySourceStorage::new(self).await?;
123
124        Ok(SequentialOrParallel::Parallel(parallel))
125    }
126
127    async fn create_sequential_source(&self) -> crate::Result<Self::SequentialSource> {
128        let seq = SequentialSafePostgresInstanceCopySourceStorage::new(self).await?;
129
130        Ok(seq)
131    }
132
133    fn supported_parallelism(&self) -> SupportedParallelism {
134        SupportedParallelism::Parallel
135    }
136}
137
138impl<'a> CopyDestinationFactory<'a> for PostgresInstanceStorage<'a> {
139    type SequentialDestination = SequentialSafePostgresInstanceCopyDestinationStorage<'a>;
140    type ParallelDestination = ParallelSafePostgresInstanceCopyDestinationStorage<'a>;
141
142    async fn create_destination(
143        &'a mut self,
144    ) -> crate::Result<SequentialOrParallel<Self::SequentialDestination, Self::ParallelDestination>>
145    {
146        let par = ParallelSafePostgresInstanceCopyDestinationStorage::new(self);
147
148        Ok(SequentialOrParallel::Parallel(par))
149    }
150
151    async fn create_sequential_destination(
152        &'a mut self,
153    ) -> crate::Result<Self::SequentialDestination> {
154        let seq = SequentialSafePostgresInstanceCopyDestinationStorage::new(self).await?;
155
156        Ok(seq)
157    }
158
159    fn supported_parallelism(&self) -> SupportedParallelism {
160        SupportedParallelism::Parallel
161    }
162}
elefant_tools/storage/postgres/postgres_instance_storage.rs

elefant_tools/storage/postgres/
postgres_instance_storage.rs
