nautilus_infrastructure/sql/
pg.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16use derive_builder::Builder;
17use regex::Regex;
18use serde::{Deserialize, Serialize};
19use sqlx::{AssertSqlSafe, ConnectOptions, PgPool, postgres::PgConnectOptions};
20
21fn validate_sql_identifier(value: &str, label: &str) -> anyhow::Result<()> {
22    if value.is_empty() {
23        anyhow::bail!("{label} must not be empty");
24    }
25
26    if !value.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') {
27        anyhow::bail!(
28            "{label} contains invalid characters (only alphanumeric and underscore allowed): {value}"
29        );
30    }
31    Ok(())
32}
33
34fn escape_sql_string(value: &str) -> String {
35    value.replace('\'', "''")
36}
37
38#[derive(Debug, Clone, Serialize, Deserialize, Builder)]
39#[serde(deny_unknown_fields)]
40#[builder(default)]
41#[cfg_attr(
42    feature = "python",
43    pyo3::pyclass(
44        module = "nautilus_trader.core.nautilus_pyo3.infrastructure",
45        from_py_object
46    )
47)]
48#[cfg_attr(
49    feature = "python",
50    pyo3_stub_gen::derive::gen_stub_pyclass(module = "nautilus_trader.infrastructure")
51)]
52#[allow(
53    clippy::unsafe_derive_deserialize,
54    reason = "config type deserializes plain field values; unsafe PyO3 methods are unrelated"
55)]
56pub struct PostgresConnectOptions {
57    pub host: String,
58    pub port: u16,
59    pub username: String,
60    pub password: String,
61    pub database: String,
62}
63
64impl PostgresConnectOptions {
65    /// Creates a new [`PostgresConnectOptions`] instance.
66    #[must_use]
67    pub const fn new(
68        host: String,
69        port: u16,
70        username: String,
71        password: String,
72        database: String,
73    ) -> Self {
74        Self {
75            host,
76            port,
77            username,
78            password,
79            database,
80        }
81    }
82
83    #[must_use]
84    pub fn connection_string(&self) -> String {
85        format!(
86            "postgres://{username}:{password}@{host}:{port}/{database}",
87            username = self.username,
88            password = self.password,
89            host = self.host,
90            port = self.port,
91            database = self.database
92        )
93    }
94
95    /// Returns the connection string with the password masked for safe logging.
96    #[must_use]
97    pub fn connection_string_masked(&self) -> String {
98        format!(
99            "postgres://{username}:***@{host}:{port}/{database}",
100            username = self.username,
101            host = self.host,
102            port = self.port,
103            database = self.database
104        )
105    }
106
107    #[must_use]
108    pub fn default_administrator() -> Self {
109        Self::new(
110            String::from("localhost"),
111            5432,
112            String::from("nautilus"),
113            String::from("pass"),
114            String::from("nautilus"),
115        )
116    }
117}
118
119impl Default for PostgresConnectOptions {
120    fn default() -> Self {
121        Self::new(
122            String::from("localhost"),
123            5432,
124            String::from("nautilus"),
125            String::from("pass"),
126            String::from("nautilus"),
127        )
128    }
129}
130
131impl From<PostgresConnectOptions> for PgConnectOptions {
132    fn from(opt: PostgresConnectOptions) -> Self {
133        Self::new()
134            .host(opt.host.as_str())
135            .port(opt.port)
136            .username(opt.username.as_str())
137            .password(opt.password.as_str())
138            .database(opt.database.as_str())
139            .disable_statement_logging()
140    }
141}
142
143/// Constructs `PostgresConnectOptions` by merging provided arguments, environment variables, and defaults.
144///
145/// # Panics
146///
147/// Panics if an environment variable for port cannot be parsed into a `u16`.
148#[must_use]
149pub fn get_postgres_connect_options(
150    host: Option<String>,
151    port: Option<u16>,
152    username: Option<String>,
153    password: Option<String>,
154    database: Option<String>,
155) -> PostgresConnectOptions {
156    let defaults = PostgresConnectOptions::default_administrator();
157    let host = host
158        .or_else(|| std::env::var("POSTGRES_HOST").ok())
159        .unwrap_or(defaults.host);
160    let port = port
161        .or_else(|| {
162            std::env::var("POSTGRES_PORT")
163                .map(|port| port.parse::<u16>().unwrap())
164                .ok()
165        })
166        .unwrap_or(defaults.port);
167    let username = username
168        .or_else(|| std::env::var("POSTGRES_USERNAME").ok())
169        .unwrap_or(defaults.username);
170    let database = database
171        .or_else(|| std::env::var("POSTGRES_DATABASE").ok())
172        .unwrap_or(defaults.database);
173    let password = password
174        .or_else(|| std::env::var("POSTGRES_PASSWORD").ok())
175        .unwrap_or(defaults.password);
176    PostgresConnectOptions::new(host, port, username, password, database)
177}
178
179/// Connects to a Postgres database with the provided connection `options` returning a connection pool.
180///
181/// # Errors
182///
183/// Returns an error if establishing the database connection fails.
184pub async fn connect_pg(options: PgConnectOptions) -> anyhow::Result<PgPool> {
185    Ok(PgPool::connect_with(options).await?)
186}
187
188/// Scans the current working directory for the `nautilus_trader` repository
189/// and constructs the path to the SQL schema directory.
190///
191/// # Errors
192///
193/// Returns an error if the `SCHEMA_DIR` environment variable is not set and the repository
194/// cannot be located in the current directory path.
195///
196/// # Panics
197///
198/// Panics if the current working directory cannot be determined or contains invalid UTF-8.
199fn get_schema_dir() -> anyhow::Result<String> {
200    std::env::var("SCHEMA_DIR").or_else(|_| {
201        let nautilus_git_repo_name = "nautilus_trader";
202        let binding = std::env::current_dir().unwrap();
203        let current_dir = binding.to_str().unwrap();
204        match current_dir.find(nautilus_git_repo_name){
205            Some(index) => {
206                let schema_path = current_dir[0..index + nautilus_git_repo_name.len()].to_string() + "/schema/sql";
207                Ok(schema_path)
208            }
209            None => anyhow::bail!("Could not calculate schema dir from current directory path or SCHEMA_DIR env variable")
210        }
211    })
212}
213
214/// Initializes the Postgres database by creating schema, roles, and executing SQL files from `schema_dir`.
215///
216/// # Errors
217///
218/// Returns an error if any SQL execution or file system operation fails.
219///
220/// # Panics
221///
222/// Panics if `schema_dir` is missing and cannot be determined or if other unwraps fail.
223#[expect(
224    clippy::too_many_lines,
225    reason = "Postgres initialization follows the ordered schema and role setup steps"
226)]
227pub async fn init_postgres(
228    pg: &PgPool,
229    database: String,
230    password: String,
231    schema_dir: Option<String>,
232) -> anyhow::Result<()> {
233    log::info!("Initializing Postgres database with target permissions and schema");
234
235    validate_sql_identifier(&database, "database")?;
236
237    // Create public schema
238    match sqlx::query("CREATE SCHEMA IF NOT EXISTS public;")
239        .execute(pg)
240        .await
241    {
242        Ok(_) => log::info!("Schema public created successfully"),
243        Err(e) => log::error!("Error creating schema public: {e:?}"),
244    }
245
246    // Create role if not exists
247    let escaped_password = escape_sql_string(&password);
248    match sqlx::query(AssertSqlSafe(format!(
249        "CREATE ROLE {database} PASSWORD '{escaped_password}' LOGIN;"
250    )))
251    .execute(pg)
252    .await
253    {
254        Ok(_) => log::info!("Role {database} created successfully"),
255        Err(e) => {
256            if e.to_string().contains("already exists") {
257                log::info!("Role {database} already exists");
258            } else {
259                log::error!("Error creating role {database}: {e:?}");
260            }
261        }
262    }
263
264    // Execute all the sql files in schema dir
265    let schema_dir = schema_dir.unwrap_or_else(|| get_schema_dir().unwrap());
266    let sql_files = vec!["types.sql", "functions.sql", "partitions.sql", "tables.sql"];
267    let plpgsql_regex =
268        Regex::new(r"\$\$ LANGUAGE plpgsql(?:[ \t\r\n]+SECURITY[ \t\r\n]+DEFINER)?;")?;
269
270    for file_name in &sql_files {
271        log::info!("Executing schema file: {file_name:?}");
272        let file_path = format!("{schema_dir}/{file_name}");
273        let sql_content = std::fs::read_to_string(&file_path)?;
274        let sql_statements: Vec<String> = match *file_name {
275            "functions.sql" | "partitions.sql" => {
276                let mut statements = Vec::new();
277                let mut last_end = 0;
278
279                for mat in plpgsql_regex.find_iter(&sql_content) {
280                    let statement = sql_content[last_end..mat.end()].to_string();
281                    if !statement.trim().is_empty() {
282                        statements.push(statement);
283                    }
284                    last_end = mat.end();
285                }
286                statements
287            }
288            _ => split_sql_statements(&sql_content),
289        };
290
291        for sql_statement in sql_statements {
292            sqlx::query(AssertSqlSafe(sql_statement.as_str()))
293                .execute(pg)
294                .await
295                .map_err(|e| {
296                    if e.to_string().contains("already exists") {
297                        log::info!("Already exists error on statement, skipping");
298                    } else {
299                        panic!("Error executing statement {sql_statement} with error: {e:?}")
300                    }
301                })
302                .unwrap();
303        }
304    }
305
306    // Grant connect
307    match sqlx::query(AssertSqlSafe(format!(
308        "GRANT CONNECT ON DATABASE {database} TO {database};"
309    )))
310    .execute(pg)
311    .await
312    {
313        Ok(_) => log::info!("Connect privileges granted to role {database}"),
314        Err(e) => log::error!("Error granting connect privileges to role {database}: {e:?}"),
315    }
316
317    // Grant all schema privileges to the role
318    match sqlx::query(AssertSqlSafe(format!(
319        "GRANT ALL PRIVILEGES ON SCHEMA public TO {database};"
320    )))
321    .execute(pg)
322    .await
323    {
324        Ok(_) => log::info!("All schema privileges granted to role {database}"),
325        Err(e) => log::error!("Error granting all privileges to role {database}: {e:?}"),
326    }
327
328    // Grant all table privileges to the role
329    match sqlx::query(AssertSqlSafe(format!(
330        "GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA public TO {database};"
331    )))
332    .execute(pg)
333    .await
334    {
335        Ok(_) => log::info!("All tables privileges granted to role {database}"),
336        Err(e) => log::error!("Error granting all privileges to role {database}: {e:?}"),
337    }
338
339    // Grant all sequence privileges to the role
340    match sqlx::query(AssertSqlSafe(format!(
341        "GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA public TO {database};"
342    )))
343    .execute(pg)
344    .await
345    {
346        Ok(_) => log::info!("All sequences privileges granted to role {database}"),
347        Err(e) => log::error!("Error granting all privileges to role {database}: {e:?}"),
348    }
349
350    // Grant all function privileges to the role
351    match sqlx::query(AssertSqlSafe(format!(
352        "GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA public TO {database};"
353    )))
354    .execute(pg)
355    .await
356    {
357        Ok(_) => log::info!("All functions privileges granted to role {database}"),
358        Err(e) => log::error!("Error granting all privileges to role {database}: {e:?}"),
359    }
360
361    Ok(())
362}
363
364// Splits semicolon-delimited SQL into individual statements.
365//
366// Skips `--` line comments and respects single-quoted string literals, so a semicolon inside a
367// comment or string literal does not split a statement. Used for the plain DDL schema files; the
368// PL/pgSQL files are split separately on their function terminators.
369fn split_sql_statements(sql: &str) -> Vec<String> {
370    let mut statements = Vec::new();
371    let mut current = String::new();
372    let mut chars = sql.chars().peekable();
373    let mut in_string = false;
374
375    while let Some(c) = chars.next() {
376        match c {
377            '\'' => {
378                // A `''` escape toggles twice, leaving the state unchanged, which is correct
379                in_string = !in_string;
380                current.push(c);
381            }
382            '-' if !in_string && chars.peek() == Some(&'-') => {
383                for next in chars.by_ref() {
384                    if next == '\n' {
385                        current.push('\n');
386                        break;
387                    }
388                }
389            }
390            ';' if !in_string => {
391                let trimmed = current.trim();
392                if !trimmed.is_empty() {
393                    statements.push(format!("{trimmed};"));
394                }
395                current.clear();
396            }
397            _ => current.push(c),
398        }
399    }
400
401    let trimmed = current.trim();
402    if !trimmed.is_empty() {
403        statements.push(format!("{trimmed};"));
404    }
405
406    statements
407}
408
409/// Drops the Postgres database with the given name using the provided connection pool.
410///
411/// # Errors
412///
413/// Returns an error if the DROP DATABASE command fails.
414pub async fn drop_postgres(pg: &PgPool, database: String) -> anyhow::Result<()> {
415    validate_sql_identifier(&database, "database")?;
416
417    // Execute drop owned
418    match sqlx::query(AssertSqlSafe(format!("DROP OWNED BY {database}")))
419        .execute(pg)
420        .await
421    {
422        Ok(_) => log::info!("Dropped owned objects by role {database}"),
423        Err(e) => {
424            let err_msg = e.to_string();
425            if err_msg.contains("2BP01") || err_msg.contains("required by the database system") {
426                log::warn!("Skipping system-required objects for role {database}");
427            } else {
428                log::error!("Error dropping owned by role {database}: {e:?}");
429            }
430        }
431    }
432
433    // Revoke connect
434    match sqlx::query(AssertSqlSafe(format!(
435        "REVOKE CONNECT ON DATABASE {database} FROM {database};"
436    )))
437    .execute(pg)
438    .await
439    {
440        Ok(_) => log::info!("Revoked connect privileges from role {database}"),
441        Err(e) => log::error!("Error revoking connect privileges from role {database}: {e:?}"),
442    }
443
444    // Revoke privileges
445    match sqlx::query(AssertSqlSafe(format!(
446        "REVOKE ALL PRIVILEGES ON DATABASE {database} FROM {database};"
447    )))
448    .execute(pg)
449    .await
450    {
451        Ok(_) => log::info!("Revoked all privileges from role {database}"),
452        Err(e) => log::error!("Error revoking all privileges from role {database}: {e:?}"),
453    }
454
455    // Execute drop schema
456    match sqlx::query("DROP SCHEMA IF EXISTS public CASCADE")
457        .execute(pg)
458        .await
459    {
460        Ok(_) => log::info!("Dropped schema public"),
461        Err(e) => log::error!("Error dropping schema public: {e:?}"),
462    }
463
464    // Drop role
465    match sqlx::query(AssertSqlSafe(format!("DROP ROLE IF EXISTS {database};")))
466        .execute(pg)
467        .await
468    {
469        Ok(_) => log::info!("Dropped role {database}"),
470        Err(e) => {
471            let err_msg = e.to_string();
472            if err_msg.contains("55006") || err_msg.contains("current user cannot be dropped") {
473                log::warn!("Cannot drop currently connected role {database}");
474            } else {
475                log::error!("Error dropping role {database}: {e:?}");
476            }
477        }
478    }
479    Ok(())
480}
481
482#[cfg(test)]
483mod tests {
484    use rstest::rstest;
485
486    use super::*;
487
488    #[rstest]
489    fn test_postgres_connect_options_toml_round_trip() {
490        let config: PostgresConnectOptions = toml::from_str(
491            r#"
492host = "localhost"
493port = 5432
494username = "nautilus"
495password = "secret"
496database = "nautilus"
497"#,
498        )
499        .unwrap();
500
501        assert_eq!(config.host, "localhost");
502        assert_eq!(config.port, 5432);
503        assert_eq!(config.username, "nautilus");
504        assert_eq!(config.database, "nautilus");
505    }
506
507    #[rstest]
508    fn test_split_sql_statements_basic() {
509        let sql = "CREATE TABLE a (id INT); CREATE TABLE b (id INT);";
510        assert_eq!(
511            split_sql_statements(sql),
512            vec!["CREATE TABLE a (id INT);", "CREATE TABLE b (id INT);"]
513        );
514    }
515
516    #[rstest]
517    fn test_split_sql_statements_ignores_semicolon_in_line_comment() {
518        // Regression: a `;` inside a `--` comment must not split the following statement
519        let sql = "\
520-- start points; a later run re-validates them.
521ALTER TABLE pool_snapshot ADD COLUMN IF NOT EXISTS validation_state TEXT;";
522        assert_eq!(
523            split_sql_statements(sql),
524            vec!["ALTER TABLE pool_snapshot ADD COLUMN IF NOT EXISTS validation_state TEXT;"]
525        );
526    }
527
528    #[rstest]
529    fn test_split_sql_statements_keeps_code_before_trailing_comment() {
530        let sql = "CREATE TABLE a (\n  id INT,  -- REFERENCES x;\n  name TEXT\n);";
531        assert_eq!(
532            split_sql_statements(sql),
533            vec!["CREATE TABLE a (\n  id INT,  \n  name TEXT\n);"]
534        );
535    }
536
537    #[rstest]
538    fn test_split_sql_statements_ignores_semicolon_in_string_literal() {
539        let sql = "INSERT INTO t VALUES ('a;b'); SELECT 1;";
540        assert_eq!(
541            split_sql_statements(sql),
542            vec!["INSERT INTO t VALUES ('a;b');", "SELECT 1;"]
543        );
544    }
545
546    #[rstest]
547    fn test_split_sql_statements_drops_comment_only_lines() {
548        let sql =
549            "------------------- ENUMS -------------------\nCREATE TYPE x AS ENUM ('A', 'B');";
550        assert_eq!(
551            split_sql_statements(sql),
552            vec!["CREATE TYPE x AS ENUM ('A', 'B');"]
553        );
554    }
555}
nautilus_infrastructure/sql/pg.rs

nautilus_infrastructure/sql/
pg.rs
