mocra-core 0.4.0

The mocra crawler framework runtime: errors, cache, utilities, domain models, downloader, data-plane queue, coordination, scheduler and engine.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88

#[cfg(feature = "store")]
use dashmap::DashMap;
#[cfg(feature = "store")]
use once_cell::sync::Lazy;
#[cfg(feature = "store")]
use sea_orm::{Database, DatabaseConnection};

// Global cache for connection pools
#[cfg(feature = "store")]
static CONNECTION_POOLS: Lazy<DashMap<String, DatabaseConnection>> = Lazy::new(DashMap::new);

pub fn create_redis_pool(
    host: &str,
    port: u16,
    db: u16,
    username: &Option<String>,
    password: &Option<String>,
    pool_size: Option<usize>,
    tls: bool,
) -> Option<deadpool_redis::Pool> {
    let addr = if tls {
        deadpool_redis::ConnectionAddr::TcpTls {
            host: host.to_string(),
            port,
            insecure: false, // Enforce certificate validation for industrial grade security
        }
    } else {
        deadpool_redis::ConnectionAddr::Tcp(host.to_string(), port)
    };

    let cfg = deadpool_redis::Config {
        connection: Some(deadpool_redis::ConnectionInfo {
            addr,
            redis: deadpool_redis::RedisConnectionInfo {
                db: db as i64,
                username: username.clone(),
                password: password.clone(),
                protocol: deadpool_redis::ProtocolVersion::RESP2,
            },
        }),
        pool: Some(deadpool_redis::PoolConfig {
            max_size: pool_size.unwrap_or(100),
            ..Default::default()
        }),
        ..Default::default()
    };
    cfg.create_pool(Some(deadpool_redis::Runtime::Tokio1)).ok()
}

#[cfg(feature = "store")]
pub async fn db_connection(
    url: Option<String>,
    schema: Option<String>,
    pool_size: Option<u32>,
    tls: Option<bool>,
) -> Option<DatabaseConnection> {
    let mut final_url = url?;
    if tls.unwrap_or(false) && !final_url.contains("sslmode=") {
        let joiner = if final_url.contains('?') { "&" } else { "?" };
        final_url = format!("{}{}sslmode=require", final_url, joiner);
    }

    let key = final_url.clone();

    if let Some(conn) = CONNECTION_POOLS.get(&key) {
        return Some(conn.clone());
    }

    let mut db_options = sea_orm::ConnectOptions::new(final_url);
    if let Some(s) = schema {
        db_options.set_schema_search_path(s);
    }
    db_options
        .max_connections(pool_size.unwrap_or(10))
        .sqlx_logging(true)
        .sqlx_logging_level(log::LevelFilter::Trace);

    match Database::connect(db_options).await {
        Ok(db) => {
            CONNECTION_POOLS.insert(key, db.clone());
            Some(db)
        }
        Err(e) => {
            log::error!("Failed to connect to database: {}", e);
            None
        }
    }
}