2ae9196
ZhenYi
feat(lib/git): add HTTP handler tree with routes and errors across analysis, archives, blobs, branches, commits, conflicts, content, diffs, hooks, lfs, logs, maintenance, packs, patches, pipelines, refs, remotes, repos, signatures, submodules, tags, trees
2ae9196
use std::sync::atomic::{AtomicUsize, Ordering};
use config::database::{DatabaseConfig, DatabaseRole as ConfigRole};
use sea_orm::{Database, DatabaseConnection, DbBackend, DbErr, ExecResult, QueryResult, Statement};
use sea_orm::prelude::async_trait::async_trait;
use crate::error::ExportError;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum DatabaseRole {
Primary,
Replica,
}
impl From<ConfigRole> for DatabaseRole {
fn from(role: ConfigRole) -> Self {
match role {
ConfigRole::Primary => DatabaseRole::Primary,
ConfigRole::Replica => DatabaseRole::Replica,
}
}
}
impl From<DatabaseRole> for ConfigRole {
fn from(role: DatabaseRole) -> Self {
match role {
DatabaseRole::Primary => ConfigRole::Primary,
DatabaseRole::Replica => ConfigRole::Replica,
}
}
}
pub struct AppDatabase {
config: DatabaseConfig,
primary: DatabaseConnection,
replicas: Vec<DatabaseConnection>,
round_robin: AtomicUsize,
}
impl AppDatabase {
pub async fn connect(config: DatabaseConfig) -> Result<Self, ExportError> {
config.validate().map_err(ExportError::InvalidConfig)?;
let primary_opts = config
.primary_connect_options()
.map_err(ExportError::InvalidConfig)?;
let primary = Database::connect(primary_opts).await?;
let mut replicas = Vec::with_capacity(config.replica_count());
for i in 0..config.replica_count() {
let opts = config
.replica_connect_options(i)
.map_err(ExportError::InvalidConfig)?;
let conn = Database::connect(opts).await?;
replicas.push(conn);
}
Ok(Self {
config,
primary,
replicas,
round_robin: AtomicUsize::new(0),
})
}
pub fn config(&self) -> &DatabaseConfig {
&self.config
}
pub fn primary(&self) -> &DatabaseConnection {
&self.primary
}
pub fn replicas(&self) -> &[DatabaseConnection] {
&self.replicas
}
pub fn replica(&self, index: usize) -> Option<&DatabaseConnection> {
self.replicas.get(index)
}
pub fn replica_count(&self) -> usize {
self.replicas.len()
}
pub fn has_replicas(&self) -> bool {
!self.replicas.is_empty()
}
pub fn write(&self) -> &DatabaseConnection {
&self.primary
}
pub fn read(&self) -> &DatabaseConnection {
self.pick_replica().unwrap_or(&self.primary)
}
pub fn for_role(&self, role: DatabaseRole) -> &DatabaseConnection {
match role {
DatabaseRole::Primary => self.write(),
DatabaseRole::Replica => self.read(),
}
}
pub fn route_for(&self, sql: &str) -> &DatabaseConnection {
if is_write_sql(sql) {
self.write()
} else {
self.read()
}
}
fn pick_replica(&self) -> Option<&DatabaseConnection> {
if self.replicas.is_empty() {
return None;
}
let idx = self.round_robin.fetch_add(1, Ordering::Relaxed) % self.replicas.len();
Some(&self.replicas[idx])
}
}
fn strip_sql_lead(sql: &str) -> &str {
let mut s = sql.trim_start();
loop {
if let Some(rest) = s.strip_prefix("--") {
let end = rest.find('\n').unwrap_or(rest.len());
s = rest[end..].trim_start();
continue;
}
if let Some(rest) = s.strip_prefix("/*") {
let end = rest.find("*/").map(|i| i + 2).unwrap_or(rest.len());
s = rest[end..].trim_start();
continue;
}
break;
}
s
}
fn first_keyword(sql: &str) -> Option<String> {
let lead = strip_sql_lead(sql);
let head: String = lead
.chars()
.take_while(|c| c.is_ascii_alphanumeric() || *c == '_')
.collect();
if head.is_empty() {
None
} else {
Some(head.to_uppercase())
}
}
fn is_write_sql(sql: &str) -> bool {
let Some(kw) = first_keyword(sql) else {
return true;
};
if kw == "WITH" {
return with_clause_is_write(sql);
}
matches!(
kw.as_str(),
"INSERT"
| "UPDATE"
| "DELETE"
| "MERGE"
| "UPSERT"
| "REPLACE"
| "TRUNCATE"
| "CREATE"
| "DROP"
| "ALTER"
| "RENAME"
| "GRANT"
| "REVOKE"
| "COMMIT"
| "ROLLBACK"
| "SAVEPOINT"
| "LOCK"
| "UNLOCK"
| "COPY"
| "CALL"
| "SET"
| "RESET"
| "DO"
)
}
fn with_clause_is_write(sql: &str) -> bool {
let upper = sql.to_uppercase();
contains_top_level(&upper, "INSERT")
|| contains_top_level(&upper, "UPDATE")
|| contains_top_level(&upper, "DELETE")
|| contains_top_level(&upper, "MERGE")
}
fn contains_top_level(haystack: &str, needle: &str) -> bool {
let bytes = haystack.as_bytes();
let n = needle.as_bytes();
let mut depth: i32 = 0;
let mut in_string = false;
let mut i = 0;
while i < bytes.len() {
let c = bytes[i] as char;
if in_string {
if c == '\'' {
if i + 1 < bytes.len() && bytes[i + 1] == b'\'' {
i += 2;
continue;
}
in_string = false;
}
i += 1;
continue;
}
match c {
'\'' => in_string = true,
'(' => depth += 1,
')' => depth -= 1,
_ if depth == 0 && i + n.len() <= bytes.len() && &bytes[i..i + n.len()] == n => {
let before_ok = i == 0
|| !(bytes[i - 1] as char).is_ascii_alphanumeric()
&& bytes[i - 1] != b'_';
let after_ok = i + n.len() == bytes.len()
|| !(bytes[i + n.len()] as char).is_ascii_alphanumeric()
&& bytes[i + n.len()] != b'_';
if before_ok && after_ok {
return true;
}
}
_ => {}
}
i += 1;
}
false
}
#[async_trait]
impl sea_orm::ConnectionTrait for AppDatabase {
fn get_database_backend(&self) -> DbBackend {
self.primary.get_database_backend()
}
async fn execute_raw(&self, stmt: Statement) -> Result<ExecResult, DbErr> {
self.route_for(&stmt.sql).execute_raw(stmt).await
}
async fn execute_unprepared(&self, sql: &str) -> Result<ExecResult, DbErr> {
self.route_for(sql).execute_unprepared(sql).await
}
async fn query_one_raw(&self, stmt: Statement) -> Result<Option<QueryResult>, DbErr> {
self.route_for(&stmt.sql).query_one_raw(stmt).await
}
async fn query_all_raw(&self, stmt: Statement) -> Result<Vec<QueryResult>, DbErr> {
self.route_for(&stmt.sql).query_all_raw(stmt).await
}
}