Lines
0 %
Functions
use std::collections::BTreeMap;
use sqlx::AssertSqlSafe;
use sqlx::MySqlConnection;
use sqlx::prelude::*;
use serde::{Deserialize, Serialize};
use crate::core::protocol::CompleteDatabaseNameResponse;
use crate::core::protocol::request_validation::GroupDenylist;
use crate::core::protocol::request_validation::validate_db_or_user_request;
use crate::core::types::DbOrUser;
use crate::core::types::MySQLDatabase;
use crate::core::types::MySQLUser;
use crate::{
core::{
common::UnixUser,
protocol::{
CreateDatabaseError, CreateDatabasesResponse, DropDatabaseError, DropDatabasesResponse,
ListAllDatabasesError, ListAllDatabasesResponse, ListDatabasesError,
ListDatabasesResponse,
},
server::{common::create_user_group_matching_regex, sql::quote_identifier},
};
const MAX_SHOW_DB_RELATED_ITEMS: usize = 5;
// NOTE: this function is unsafe because it does no input validation.
pub(super) async fn unsafe_database_exists(
database_name: &str,
connection: &mut MySqlConnection,
) -> Result<bool, sqlx::Error> {
let result =
sqlx::query("SELECT SCHEMA_NAME FROM information_schema.SCHEMATA WHERE SCHEMA_NAME = ?")
.bind(database_name)
.fetch_optional(connection)
.await;
if let Err(err) = &result {
tracing::error!(
"Failed to check if database '{}' exists: {:?}",
&database_name,
err
);
}
Ok(result?.is_some())
pub async fn complete_database_name(
database_prefix: &str,
unix_user: &UnixUser,
_db_is_mariadb: bool,
group_denylist: &GroupDenylist,
) -> CompleteDatabaseNameResponse {
let result = sqlx::query(
r"
SELECT CAST(`SCHEMA_NAME` AS CHAR(64)) AS `database`
FROM `information_schema`.`SCHEMATA`
WHERE `SCHEMA_NAME` NOT IN ('information_schema', 'performance_schema', 'mysql', 'sys')
AND `SCHEMA_NAME` REGEXP ?
AND `SCHEMA_NAME` LIKE ?
",
)
.bind(create_user_group_matching_regex(unix_user, group_denylist))
.bind(format!("{database_prefix}%"))
.fetch_all(connection)
match result {
Ok(rows) => rows
.into_iter()
.filter_map(|row| {
let database: String = row.try_get("database").ok()?;
Some(database.into())
})
.collect(),
Err(err) => {
"Failed to complete database name for prefix '{}' and user '{}': {:?}",
database_prefix,
unix_user.username,
vec![]
pub async fn create_databases(
database_names: &[MySQLDatabase],
) -> CreateDatabasesResponse {
let mut results = BTreeMap::new();
for database_name in database_names.iter().cloned() {
if let Err(err) = validate_db_or_user_request(
&DbOrUser::Database(database_name.clone()),
unix_user,
group_denylist,
.map_err(CreateDatabaseError::ValidationError)
{
results.insert(database_name.clone(), Err(err));
continue;
match unsafe_database_exists(&database_name, &mut *connection).await {
Ok(true) => {
results.insert(
database_name.clone(),
Err(CreateDatabaseError::DatabaseAlreadyExists),
Err(CreateDatabaseError::MySqlError(err.to_string())),
_ => {}
let statement = AssertSqlSafe(format!(
"CREATE DATABASE {}",
quote_identifier(&database_name)
));
let result = sqlx::query(statement)
.execute(&mut *connection)
.await
.map(|_| ())
.map_err(|err| CreateDatabaseError::MySqlError(err.to_string()));
tracing::error!("Failed to create database '{}': {:?}", &database_name, err);
results.insert(database_name, result);
results
pub async fn drop_databases(
) -> DropDatabasesResponse {
.map_err(DropDatabaseError::ValidationError)
Ok(false) => {
Err(DropDatabaseError::DatabaseDoesNotExist),
Err(DropDatabaseError::MySqlError(err.to_string())),
"DROP DATABASE {}",
.map_err(|err| DropDatabaseError::MySqlError(err.to_string()));
tracing::error!("Failed to drop database '{}': {:?}", &database_name, err);
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct DatabaseRow {
pub database: MySQLDatabase,
pub tables: Vec<String>,
pub table_count: u64,
pub users: Vec<MySQLUser>,
pub user_count: u64,
pub collation: Option<String>,
pub character_set: Option<String>,
pub size_bytes: u64,
impl FromRow<'_, sqlx::mysql::MySqlRow> for DatabaseRow {
fn from_row(row: &sqlx::mysql::MySqlRow) -> Result<Self, sqlx::Error> {
Ok(DatabaseRow {
database: row.try_get::<String, _>("database")?.into(),
tables: {
let s: Option<String> = row.try_get("tables")?;
s.and_then(|s| {
if s.is_empty() {
None
} else {
Some(s.split(',').map(std::borrow::ToOwned::to_owned).collect())
.unwrap_or_default()
table_count: row.try_get::<u64, _>("table_count")?,
users: {
let s: Option<String> = row.try_get("users")?;
Some(s.split(',').map(|s| s.to_owned().into()).collect())
user_count: row.try_get::<u64, _>("user_count")?,
collation: row.try_get::<Option<String>, _>("collation")?,
character_set: row.try_get::<Option<String>, _>("character_set")?,
size_bytes: row.try_get::<u64, _>("size_bytes")?,
fn list_database_query(include_all_tables_and_users: bool) -> AssertSqlSafe<String> {
let limit_clause = if include_all_tables_and_users {
"".to_string()
format!(" LIMIT {}", MAX_SHOW_DB_RELATED_ITEMS)
AssertSqlSafe(format!(
SELECT
CAST(s.SCHEMA_NAME AS CHAR(64)) AS `database`,
t.tables,
CAST(COALESCE(sz.table_count, 0) AS UNSIGNED) AS table_count,
u.users,
CAST(COALESCE(uc.user_count, 0) AS UNSIGNED) AS user_count,
s.DEFAULT_COLLATION_NAME AS `collation`,
s.DEFAULT_CHARACTER_SET_NAME AS `character_set`,
CAST(COALESCE(sz.size_bytes, 0) AS UNSIGNED) AS size_bytes
FROM information_schema.SCHEMATA s
LEFT JOIN (
x.TABLE_SCHEMA,
GROUP_CONCAT(x.TABLE_NAME ORDER BY x.TABLE_NAME SEPARATOR ',') AS tables
FROM (
TABLE_SCHEMA,
TABLE_NAME
FROM information_schema.TABLES
WHERE TABLE_SCHEMA = ?
ORDER BY TABLE_NAME{limit_clause}
) x
GROUP BY x.TABLE_SCHEMA
) t
ON t.TABLE_SCHEMA = s.SCHEMA_NAME
x.DB,
GROUP_CONCAT(x.User ORDER BY x.User SEPARATOR ',') AS users
SELECT DISTINCT
DB,
User
FROM mysql.db
WHERE DB = ?
ORDER BY User{limit_clause}
GROUP BY x.DB
) u
ON u.DB = s.SCHEMA_NAME
SUM(DATA_LENGTH + INDEX_LENGTH) AS size_bytes,
COUNT(*) AS table_count
GROUP BY TABLE_SCHEMA
) sz
ON sz.TABLE_SCHEMA = s.SCHEMA_NAME
COUNT(DISTINCT User) AS user_count
GROUP BY DB
) uc
ON uc.DB = s.SCHEMA_NAME
WHERE s.SCHEMA_NAME REGEXP ?
AND s.SCHEMA_NAME NOT IN (
'information_schema',
'performance_schema',
'mysql',
'sys'
"
))
pub async fn list_databases(
include_all_tables_and_users: bool,
) -> ListDatabasesResponse {
.map_err(ListDatabasesError::ValidationError)
let query = list_database_query(include_all_tables_and_users);
let result = sqlx::query_as::<_, DatabaseRow>(query)
.bind(database_name.to_string())
.fetch_optional(&mut *connection)
.map_err(|err| ListDatabasesError::MySqlError(err.to_string()))
.and_then(|database| {
database.map_or_else(|| Err(ListDatabasesError::DatabaseDoesNotExist), Ok)
});
tracing::error!("Failed to list database '{}': {:?}", &database_name, err);
// TODO: should we assert that the users are also owned by the unix_user from the request?
fn list_all_databases_for_user_query(include_all_tables_and_users: bool) -> AssertSqlSafe<String> {
let row_limit_clause = if include_all_tables_and_users {
String::new()
format!("WHERE row_num <= {MAX_SHOW_DB_RELATED_ITEMS}")
TABLE_NAME,
ROW_NUMBER() OVER (PARTITION BY TABLE_SCHEMA ORDER BY TABLE_NAME) AS row_num
WHERE TABLE_SCHEMA REGEXP ?
{row_limit_clause}
User,
ROW_NUMBER() OVER (PARTITION BY DB ORDER BY User) AS row_num
SELECT DISTINCT DB, User
WHERE DB REGEXP ?
) d
ORDER BY s.SCHEMA_NAME
pub async fn list_all_databases_for_user(
) -> ListAllDatabasesResponse {
let query = list_all_databases_for_user_query(include_all_tables_and_users);
let user_group_regex = create_user_group_matching_regex(unix_user, group_denylist);
.bind(&user_group_regex)
.map_err(|err| ListAllDatabasesError::MySqlError(err.to_string()));
"Failed to list databases for user '{}': {:?}",
result