use std::sync::Arc;
use std::time::{Duration, Instant};
#[cfg(feature = "permission")]
use crate::access::permission::{PermissionAction, PermissionContext};
#[cfg(feature = "sql-parser")]
use crate::access::security::{DdlGuard, DdlValidationResult};
#[cfg(all(feature = "sql-parser", feature = "permission"))]
use crate::access::sql_parser::SqlParser;
#[cfg(feature = "sql-parser")]
use crate::access::sql_parser::is_ddl_operation;
use crate::database::pool::db_pool::{DatabaseConnection, DbConnection, DbPool, DbPoolInner};
use crate::foundation::error::{DbError, DbResult};
#[cfg(feature = "metrics")]
use crate::observability::metrics::MetricsCollector;
use async_trait::async_trait;
#[cfg(not(any(feature = "permission", feature = "sql-parser")))]
compile_error!("Either 'permission' or 'sql-parser' feature must be enabled for Session functionality");
use sea_orm::{ConnectionTrait, DatabaseTransaction, ExecResult, TransactionTrait};
use tokio::sync::Mutex;
struct SessionState {
transaction: Option<Arc<DatabaseTransaction>>,
last_write: Option<Instant>,
}
pub struct Session {
connection: Option<DbConnection>,
pool: Arc<DbPool>,
pool_inner: Arc<DbPoolInner>,
role: String,
#[cfg(feature = "permission")]
permission_ctx: PermissionContext,
state: Mutex<SessionState>,
#[cfg(feature = "metrics")]
metrics_collector: Option<Arc<MetricsCollector>>,
}
impl Session {
pub(crate) fn new(connection: DbConnection, pool: Arc<DbPool>, pool_inner: Arc<DbPoolInner>, role: String) -> Self {
#[cfg(feature = "permission")]
let permission_ctx = PermissionContext::new(role.clone(), pool_inner.policy_cache.clone());
#[cfg(feature = "metrics")]
let metrics = pool_inner.metrics_collector.clone();
Session {
connection: Some(connection),
pool,
pool_inner,
role,
#[cfg(feature = "permission")]
permission_ctx,
state: Mutex::new(SessionState {
transaction: None,
last_write: None,
}),
#[cfg(feature = "metrics")]
metrics_collector: metrics,
}
}
pub fn role(&self) -> &str {
&self.role
}
#[cfg(feature = "permission")]
pub fn permission_ctx(&self) -> &PermissionContext {
&self.permission_ctx
}
pub async fn mark_write(&self) {
let mut state = self.state.lock().await;
state.last_write = Some(Instant::now());
}
#[cfg(feature = "permission")]
pub async fn check_permission(&self, table: &str, operation: &PermissionAction) -> Result<(), DbError> {
if self.role == self.pool_inner.admin_role {
return Ok(());
}
if self.permission_ctx.check_table_access(table, operation).await {
Ok(())
} else {
Err(permission_denied(operation, table))
}
}
pub async fn is_in_transaction(&self) -> bool {
let state = self.state.lock().await;
state.transaction.is_some()
}
pub async fn begin_transaction(&self) -> Result<(), DbError> {
{
let state = self.state.lock().await;
if state.transaction.is_some() {
return Err(DbError::Transaction("Already in transaction".to_string()));
}
}
let conn = self.connection()?;
let transaction = conn
.begin()
.await
.map_err(|e| DbError::Transaction(format!("Failed to begin transaction: {}", e)))?;
let mut state = self.state.lock().await;
if state.transaction.is_some() {
let _ = transaction.rollback().await;
return Err(DbError::Transaction(
"Already in transaction (concurrent begin detected)".to_string(),
));
}
state.transaction = Some(Arc::new(transaction));
Ok(())
}
pub async fn commit(&self) -> Result<(), DbError> {
let transaction_arc = {
let mut state = self.state.lock().await;
state
.transaction
.take()
.ok_or_else(|| DbError::Transaction("No active transaction to commit".to_string()))?
};
let transaction = Arc::try_unwrap(transaction_arc).map_err(|_| {
DbError::Transaction("Cannot commit: transaction is in use by a concurrent query".to_string())
})?;
transaction
.commit()
.await
.map_err(|e| DbError::Transaction(e.to_string()))?;
let mut state = self.state.lock().await;
state.last_write = None;
Ok(())
}
pub async fn rollback(&self) -> Result<(), DbError> {
let transaction_arc = {
let mut state = self.state.lock().await;
if state.transaction.is_none() {
return Err(DbError::Transaction("Not in transaction".to_string()));
}
state
.transaction
.take()
.ok_or_else(|| DbError::Transaction("No active transaction to rollback".to_string()))?
};
let transaction = Arc::try_unwrap(transaction_arc).map_err(|_| {
DbError::Transaction("Cannot rollback: transaction is in use by a concurrent query".to_string())
})?;
transaction
.rollback()
.await
.map_err(|e| DbError::Transaction(format!("Failed to rollback transaction: {}", e)))?;
Ok(())
}
pub async fn should_use_master(&self) -> bool {
let state = self.state.lock().await;
if state.transaction.is_some() {
return true;
}
state
.last_write
.map(|t| t.elapsed() < Duration::from_secs(5))
.unwrap_or(false)
}
pub fn connection(&self) -> Result<&DatabaseConnection, DbError> {
self.connection
.as_ref()
.ok_or_else(|| DbError::Config("Connection not available - Session may have been invalidated".to_string()))?
.as_sea_orm()
}
#[allow(dead_code)]
#[cfg(feature = "migration")]
pub fn create_migration_executor(
&self,
db_type: crate::foundation::config::DatabaseType,
) -> Result<super::MigrationExecutor, DbError> {
let conn = self.connection()?.clone();
Ok(super::MigrationExecutor::new(conn, db_type))
}
#[cfg_attr(feature = "tracing", tracing::instrument(skip(self), fields(db.role = %self.role)))]
pub async fn execute_raw(&self, sql: &str) -> DbResult<ExecResult> {
#[cfg(feature = "sql-parser")]
{
if is_ddl_operation(sql) {
return Err(DbError::Permission(
"DDL operations are not allowed in this context".to_string(),
));
}
}
#[cfg(not(feature = "sql-parser"))]
{
let _ = sql;
return Err(DbError::Permission(
"execute_raw requires the sql-parser feature to be enabled".to_string(),
));
}
#[cfg(feature = "sql-parser")]
{
#[cfg(all(feature = "sql-parser", feature = "permission"))]
{
let parser = SqlParser::shared().await;
match parser.parse_operation_async(sql).await {
Ok(Some((table_name, action))) => {
if table_name.is_empty() || is_invalid_table_name(&table_name) {
return Err(DbError::Permission(
"Failed to extract table name for permission checking".to_string(),
));
}
if self.role == self.pool_inner.admin_role {
} else if !self.permission_ctx.check_table_access(&table_name, &action).await {
return Err(permission_denied(&action, &table_name));
}
}
Ok(None) => {
return Err(DbError::Permission(
"SQL statement requires a valid table name for permission checking".to_string(),
));
}
Err(_) => {
return Err(DbError::Permission(
"Failed to parse SQL statement for permission checking".to_string(),
));
}
}
}
let tx_opt: Option<Arc<DatabaseTransaction>> = {
let state = self.state.lock().await;
state.transaction.clone()
};
if let Some(tx) = tx_opt {
return tx.execute_unprepared(sql).await.map_err(DbError::Connection);
}
let conn = self.connection()?;
conn.execute_unprepared(sql).await.map_err(DbError::Connection)
}
}
pub async fn execute_raw_ddl(&self, sql: &str) -> DbResult<ExecResult> {
if self.role != self.pool_inner.admin_role {
return Err(DbError::Permission(format!(
"DDL operations are only allowed for admin role. Current role: '{}', Admin role: '{}'",
self.role, self.pool_inner.admin_role
)));
}
#[cfg(feature = "sql-parser")]
{
let guard = DdlGuard::new();
match guard.validate(sql) {
Ok(DdlValidationResult::Allowed) => {
}
Ok(DdlValidationResult::Forbidden(reason)) => {
return Err(DbError::Permission(format!("DDL operation not allowed: {}", reason)));
}
Ok(DdlValidationResult::ParseError(error)) => {
return Err(DbError::Config(format!("Failed to parse DDL SQL: {}", error)));
}
Err(error) => {
return Err(DbError::Config(format!("DDL validation error: {}", error)));
}
}
}
let conn = self.connection()?;
conn.execute_unprepared(sql).await.map_err(DbError::Connection)
}
#[cfg(feature = "duckdb")]
pub async fn execute_duckdb(&self, sql: &str) -> DbResult<Vec<crate::database::pool::DuckDbRow>> {
#[cfg(feature = "sql-parser")]
{
if is_ddl_operation(sql) {
return Err(DbError::Permission(
"DDL operations are not allowed in DuckDB query context".to_string(),
));
}
}
#[cfg(not(feature = "sql-parser"))]
{
let _ = sql;
return Err(DbError::Permission(
"execute_duckdb requires the sql-parser feature to be enabled for security checks".to_string(),
));
}
#[cfg(feature = "sql-parser")]
{
#[cfg(all(feature = "sql-parser", feature = "permission"))]
{
let parser = SqlParser::shared().await;
match parser.parse_operation_async(sql).await {
Ok(Some((table_name, action))) => {
if table_name.is_empty() || is_invalid_table_name(&table_name) {
return Err(DbError::Permission(
"Failed to extract table name for permission checking".to_string(),
));
}
if self.role != self.pool_inner.admin_role
&& !self.permission_ctx.check_table_access(&table_name, &action).await
{
return Err(permission_denied(&action, &table_name));
}
}
Ok(None) => {
if self.role != self.pool_inner.admin_role {
return Err(DbError::Permission(
"SQL statement requires a valid table name for permission checking".to_string(),
));
}
}
Err(_) => {
return Err(DbError::Permission(
"Failed to parse SQL statement for permission checking".to_string(),
));
}
}
}
let conn = self
.connection
.as_ref()
.ok_or_else(|| DbError::Config("Connection not available".to_string()))?;
let duck_conn = conn.as_duckdb()?;
duck_conn.query(sql).await
}
}
#[cfg(feature = "duckdb")]
pub async fn execute_duckdb_raw(&self, sql: &str) -> DbResult<crate::database::pool::DuckDbExecResult> {
#[cfg(feature = "sql-parser")]
{
if is_ddl_operation(sql) {
if self.role == self.pool_inner.admin_role {
let guard = DdlGuard::new();
match guard.validate(sql) {
Ok(DdlValidationResult::Allowed) => {
let conn = self
.connection
.as_ref()
.ok_or_else(|| DbError::Config("Connection not available".to_string()))?;
let duck_conn = conn.as_duckdb()?;
return duck_conn.execute(sql).await;
}
Ok(DdlValidationResult::Forbidden(reason)) => {
return Err(DbError::Permission(format!("DDL operation not allowed: {}", reason)));
}
Ok(DdlValidationResult::ParseError(error)) => {
return Err(DbError::Config(format!("Failed to parse DDL SQL: {}", error)));
}
Err(error) => {
return Err(DbError::Config(format!("DDL validation error: {}", error)));
}
}
} else {
return Err(DbError::Permission(format!(
"DDL operations are only allowed for admin role in DuckDB context. Current role: '{}', Admin role: '{}'",
self.role, self.pool_inner.admin_role
)));
}
}
}
#[cfg(not(feature = "sql-parser"))]
{
let _ = sql;
return Err(DbError::Permission(
"execute_duckdb_raw requires the sql-parser feature to be enabled for security checks".to_string(),
));
}
#[cfg(feature = "sql-parser")]
{
#[cfg(all(feature = "sql-parser", feature = "permission"))]
{
let parser = SqlParser::shared().await;
match parser.parse_operation_async(sql).await {
Ok(Some((table_name, action))) => {
if table_name.is_empty() || is_invalid_table_name(&table_name) {
return Err(DbError::Permission(
"Failed to extract table name for permission checking".to_string(),
));
}
if self.role != self.pool_inner.admin_role
&& !self.permission_ctx.check_table_access(&table_name, &action).await
{
return Err(permission_denied(&action, &table_name));
}
}
Ok(None) => {
if self.role != self.pool_inner.admin_role {
return Err(DbError::Permission(
"SQL statement requires a valid table name for permission checking".to_string(),
));
}
}
Err(_) => {
return Err(DbError::Permission(
"Failed to parse SQL statement for permission checking".to_string(),
));
}
}
}
let conn = self
.connection
.as_ref()
.ok_or_else(|| DbError::Config("Connection not available".to_string()))?;
let duck_conn = conn.as_duckdb()?;
duck_conn.execute(sql).await
}
}
#[cfg_attr(feature = "tracing", tracing::instrument(skip(self), fields(db.role = %self.role)))]
pub async fn execute(&self, sql: &str) -> DbResult<ExecResult> {
let start = Instant::now();
#[cfg(feature = "sql-parser")]
check_ddl_operation(sql)?;
#[cfg(feature = "permission")]
{
let parsed = parse_sql_for_permission(sql).await?;
match parsed {
Some((table_name, action)) => {
if table_name.is_empty() || is_invalid_table_name(&table_name) {
return Err(DbError::Permission(
"Failed to extract table name for permission checking".to_string(),
));
}
self.check_permission(&table_name, &action).await?;
let result = self.execute_raw(sql).await?;
self.record_metrics_and_mark_write(&action, start).await;
Ok(result)
}
None => {
let result = self.execute_raw(sql).await?;
Ok(result)
}
}
}
#[cfg(not(feature = "permission"))]
{
let result = self.execute_raw(sql).await?;
Ok(result)
}
}
#[cfg(feature = "permission")]
#[cfg_attr(feature = "tracing", tracing::instrument(skip(self, operation), fields(db.role = %self.role)))]
pub async fn execute_with_operation(&self, sql: &str, operation: &PermissionAction) -> DbResult<ExecResult> {
let start = Instant::now();
#[cfg(feature = "sql-parser")]
{
if is_ddl_operation(sql) {
return Err(DbError::Permission(
"DDL operations are not allowed in this context".to_string(),
));
}
}
let table_name = extract_table_name(sql);
#[cfg(feature = "permission")]
{
if !table_name.is_empty() && !self.permission_ctx.check_table_access(&table_name, operation).await {
return Err(permission_denied(operation, &table_name));
}
}
let result = self.execute_raw(sql).await?;
let duration = start.elapsed();
self.record_query_metrics(&format!("{:?}", operation), duration, true);
#[cfg(feature = "permission")]
{
if is_write_action(operation) {
self.mark_write().await;
}
}
Ok(result)
}
pub async fn batch_execute(&self, sqls: Vec<&str>) -> DbResult<Vec<DbResult<ExecResult>>> {
let mut results = Vec::new();
for sql in sqls {
let result = self.execute(sql).await;
results.push(result);
}
Ok(results)
}
pub async fn batch_execute_in_transaction(&self, sqls: Vec<&str>) -> DbResult<Vec<ExecResult>> {
self.begin_transaction().await?;
let mut results = Vec::new();
let mut last_error = None;
for sql in sqls {
match self.execute_raw(sql).await {
Ok(result) => results.push(result),
Err(e) => {
last_error = Some(e);
break;
}
}
}
if let Some(error) = last_error {
self.rollback().await?;
Err(error)
} else {
self.commit().await?;
Ok(results)
}
}
#[cfg(feature = "metrics")]
fn record_query_metrics(&self, query_type: &str, duration: Duration, success: bool) {
if let Some(metrics) = &self.metrics_collector {
metrics.record_query(query_type, duration, success, None);
}
}
#[cfg(not(feature = "metrics"))]
fn record_query_metrics(&self, _query_type: &str, _duration: Duration, _success: bool) {
}
#[cfg(feature = "permission")]
async fn record_metrics_and_mark_write(&self, action: &PermissionAction, start: Instant) {
let duration = start.elapsed();
self.record_query_metrics(&format!("{:?}", action), duration, true);
if is_write_action(action) {
self.mark_write().await;
}
}
pub async fn check_table_permission(&self, _table_name: &str, _operation: &str) -> DbResult<()> {
#[cfg(feature = "permission")]
{
let action = match _operation {
"INSERT" => PermissionAction::Insert,
"SELECT" => PermissionAction::Select,
"UPDATE" => PermissionAction::Update,
"DELETE" => PermissionAction::Delete,
_ => return Err(DbError::Permission(format!("Unknown operation: {}", _operation))),
};
if self.role == self.pool_inner.admin_role {
} else if !self.permission_ctx.check_table_access(_table_name, &action).await {
return Err(permission_denied(_operation, _table_name));
}
}
Ok(())
}
#[cfg(feature = "metrics")]
pub fn record_metric(&self, operation: &str, table_name: &str, success: bool) {
if let Some(metrics) = &self.metrics_collector {
let bytes = Some(table_name.len() as u64);
metrics.record_query(operation, std::time::Duration::from_millis(0), success, bytes);
}
}
}
#[cfg(feature = "permission")]
fn is_invalid_table_name(table_name: &str) -> bool {
let table_name = table_name.trim();
if table_name.is_empty() {
return true;
}
for part in table_name.split('.') {
let part = part.trim();
if part.is_empty() {
return true;
}
let unquoted = part
.strip_prefix('"')
.and_then(|s| s.strip_suffix('"'))
.or_else(|| part.strip_prefix('`').and_then(|s| s.strip_suffix('`')))
.or_else(|| part.strip_prefix('\'').and_then(|s| s.strip_suffix('\'')))
.unwrap_or(part)
.trim();
if unquoted.is_empty() {
return true;
}
}
false
}
#[cfg(feature = "permission")]
fn permission_denied(action: &(impl std::fmt::Display + ?Sized), table: &(impl std::fmt::Display + ?Sized)) -> DbError {
DbError::Permission(format!("Permission denied for {} on {}", action, table))
}
#[cfg(feature = "permission")]
fn is_write_action(action: &PermissionAction) -> bool {
matches!(
action,
PermissionAction::Insert | PermissionAction::Update | PermissionAction::Delete
)
}
#[cfg(feature = "sql-parser")]
fn check_ddl_operation(sql: &str) -> DbResult<()> {
if is_ddl_operation(sql) {
return Err(DbError::Permission(
"DDL operations are not allowed in this context".to_string(),
));
}
Ok(())
}
impl Drop for Session {
fn drop(&mut self) {
if let Some(conn) = self.connection.take() {
self.pool.release_connection(conn);
}
}
}
#[cfg(feature = "permission")]
fn extract_table_name(sql: &str) -> String {
let sql_upper = sql.to_uppercase();
if sql_upper.contains("FROM ") {
if let Some(start) = sql_upper.find("FROM ") {
let rest = &sql[start + 5..];
if let Some(end) = rest.find(|c| [' ', ',', ';', '(', ')'].contains(&c)) {
return rest[..end].trim().to_string();
} else {
return rest.trim().to_string();
}
}
}
if sql_upper.contains("INTO ") {
if let Some(start) = sql_upper.find("INTO ") {
let rest = &sql[start + 5..];
if let Some(end) = rest.find(|c| [' ', '(', ';'].contains(&c)) {
return rest[..end].trim().to_string();
} else {
return rest.trim().to_string();
}
}
}
if sql_upper.contains("UPDATE ") {
if let Some(start) = sql_upper.find("UPDATE ") {
let rest = &sql[start + 7..];
if let Some(end) = rest.find(|c| [' ', ';'].contains(&c)) {
return rest[..end].trim().to_string();
} else {
return rest.trim().to_string();
}
}
}
String::new()
}
#[cfg(all(feature = "permission", not(feature = "sql-parser")))]
fn parse_table_and_action(sql: &str) -> (String, PermissionAction) {
let table_name = extract_table_name(sql);
let sql_upper = sql.trim_start().to_uppercase();
let action = if sql_upper.starts_with("INSERT") {
PermissionAction::Insert
} else if sql_upper.starts_with("UPDATE") {
PermissionAction::Update
} else if sql_upper.starts_with("DELETE") {
PermissionAction::Delete
} else {
PermissionAction::Select
};
(table_name, action)
}
#[cfg(feature = "permission")]
async fn parse_sql_for_permission(sql: &str) -> DbResult<Option<(String, PermissionAction)>> {
#[cfg(feature = "sql-parser")]
{
let parser = SqlParser::shared().await;
match parser.parse_operation_async(sql).await {
Ok(Some((table, action))) => Ok(Some((table, action))),
Ok(None) => Ok(None),
Err(_) => Ok(None),
}
}
#[cfg(not(feature = "sql-parser"))]
{
Ok(Some(parse_table_and_action(sql)))
}
}
#[async_trait]
impl super::DatabaseSession for Session {
async fn execute(&self, sql: &str) -> crate::DbResult<ExecResult> {
Ok(self.execute(sql).await?)
}
async fn execute_raw(&self, sql: &str) -> crate::DbResult<ExecResult> {
Ok(self.execute_raw(sql).await?)
}
async fn execute_raw_ddl(&self, sql: &str) -> crate::DbResult<ExecResult> {
Ok(self.execute_raw_ddl(sql).await?)
}
async fn begin_transaction(&self) -> crate::DbResult<()> {
Ok(self.begin_transaction().await?)
}
async fn commit(&self) -> crate::DbResult<()> {
Ok(self.commit().await?)
}
async fn rollback(&self) -> crate::DbResult<()> {
Ok(self.rollback().await?)
}
fn role(&self) -> &str {
self.role()
}
async fn is_in_transaction(&self) -> bool {
self.is_in_transaction().await
}
}