mirror of
https://github.com/t8y2/dbx.git
synced 2026-10-19 16:16:45 +00:00
fix(plugins): wire plugin commands and restore build
This commit is contained in:
@@ -16,8 +16,6 @@ use crate::connection_secrets::{
|
||||
CASSANDRA_KEYSTORE_PASSWORD_KEY, CASSANDRA_TRUSTSTORE_PASSWORD_KEY, MQ_AUTH_API_KEY_VALUE_KEY,
|
||||
MQ_AUTH_CLIENT_SECRET_KEY, MQ_AUTH_PASSWORD_KEY, MQ_AUTH_TOKEN_KEY, MQ_TOKEN_SIGNING_KEY, NACOS_AUTH_PASSWORD_KEY,
|
||||
NACOS_RNACOS_CONSOLE_PASSWORD_KEY,
|
||||
MQ_AUTH_API_KEY_VALUE_KEY, MQ_AUTH_CLIENT_SECRET_KEY, MQ_AUTH_PASSWORD_KEY, MQ_AUTH_TOKEN_KEY,
|
||||
MQ_TOKEN_SIGNING_KEY, NACOS_AUTH_PASSWORD_KEY, NACOS_RNACOS_CONSOLE_PASSWORD_KEY,
|
||||
};
|
||||
use crate::models::connection::{ConnectionConfig, DatabaseType, TransportLayerConfig};
|
||||
use crate::saved_sql::SavedSqlLibrary;
|
||||
@@ -1658,7 +1656,6 @@ mod tests {
|
||||
CASSANDRA_KEYSTORE_PASSWORD_KEY, CASSANDRA_TRUSTSTORE_PASSWORD_KEY, NACOS_AUTH_PASSWORD_KEY,
|
||||
NACOS_RNACOS_CONSOLE_PASSWORD_KEY,
|
||||
};
|
||||
use crate::connection_secrets::{NACOS_AUTH_PASSWORD_KEY, NACOS_RNACOS_CONSOLE_PASSWORD_KEY};
|
||||
use crate::models::connection::{
|
||||
default_redis_key_separator, ConnectionConfig, DatabaseType, SshTunnelConfig, TransportLayerConfig,
|
||||
};
|
||||
|
||||
@@ -29,7 +29,7 @@ use crate::models::connection::{
|
||||
use crate::mongo_oidc::MongoOidcBrowserOpener;
|
||||
use crate::nacos::config::{NACOS_CONSOLE_SESSION_PASSWORD, NACOS_PRIMARY_SESSION_PASSWORD};
|
||||
use crate::path_utils::expand_tilde;
|
||||
use crate::plugins::{PluginDriverSession, PluginRegistry, PluginRuntimeEnv};
|
||||
use crate::plugins::{PluginConnectionHandle, PluginDriverSession, PluginHost, PluginRegistry, PluginRuntimeEnv};
|
||||
use crate::query_cancel::RunningQueries;
|
||||
use crate::session_credentials::SessionCredentialStore;
|
||||
use crate::storage::{normalize_duckdb_worker_max_processes, Storage, DUCKDB_WORKER_MAX_PROCESSES_DEFAULT};
|
||||
@@ -110,6 +110,7 @@ pub enum PoolKind {
|
||||
config: Arc<ConnectionConfig>,
|
||||
session: Arc<PluginDriverSession>,
|
||||
},
|
||||
PluginConnection(PluginConnectionHandle),
|
||||
/// Message queue admin connection (not a data query pool; serves as a
|
||||
/// marker that this connection_id is a valid MQ admin connection).
|
||||
MessageQueue,
|
||||
@@ -331,6 +332,7 @@ pub struct AppState {
|
||||
pub http_tunnels: HttpTunnelManager,
|
||||
pub storage: Storage,
|
||||
pub plugins: PluginRegistry,
|
||||
pub plugin_host: PluginHost,
|
||||
pub agent_manager: crate::agent_manager::AgentManager,
|
||||
pub nacos_registry: crate::nacos::NacosAdminRegistry,
|
||||
duckdb_worker_process_isolation: AtomicBool,
|
||||
@@ -1357,7 +1359,8 @@ impl AppState {
|
||||
proxy_tunnels: ProxyTunnelManager::new(),
|
||||
http_tunnels: HttpTunnelManager::new(),
|
||||
storage,
|
||||
plugins: PluginRegistry::new(plugin_dir),
|
||||
plugins: PluginRegistry::new(plugin_dir.clone()),
|
||||
plugin_host: PluginHost::new(PluginRegistry::new(plugin_dir)),
|
||||
agent_manager: crate::agent_manager::AgentManager::new_with_base_dir_and_app_version(
|
||||
agent_dir,
|
||||
app_version,
|
||||
@@ -2777,6 +2780,7 @@ impl AppState {
|
||||
}
|
||||
self.external_driver_pool("jdbc", &jdbc_config).await?
|
||||
}
|
||||
DatabaseType::Plugin => return Err("Plugin-owned connections use the plugin host".to_string()),
|
||||
#[cfg(feature = "mq-admin")]
|
||||
DatabaseType::MessageQueue => {
|
||||
// MQ admin connections don't hold a data query pool. We just test
|
||||
@@ -3830,6 +3834,7 @@ impl AppState {
|
||||
PoolKind::Sqlite(_)
|
||||
| PoolKind::DuckDbWorker(_)
|
||||
| PoolKind::ExternalDriver { .. }
|
||||
| PoolKind::PluginConnection(_)
|
||||
| PoolKind::MessageQueue
|
||||
| PoolKind::Nacos
|
||||
| PoolKind::Consul(_) => false,
|
||||
@@ -4864,6 +4869,7 @@ impl AppState {
|
||||
PoolKind::Sqlite(_)
|
||||
| PoolKind::DuckDbWorker(_)
|
||||
| PoolKind::ExternalDriver { .. }
|
||||
| PoolKind::PluginConnection(_)
|
||||
| PoolKind::MessageQueue
|
||||
| PoolKind::Nacos
|
||||
| PoolKind::Consul(_) => true,
|
||||
@@ -4876,7 +4882,6 @@ impl AppState {
|
||||
false
|
||||
}
|
||||
},
|
||||
PoolKind::Redis(_) => unreachable!("Redis handled separately"),
|
||||
};
|
||||
if !healthy && !matches!(pool, PoolKind::Agent(_)) {
|
||||
dead_pools.push((key.clone(), checked.publication.clone()));
|
||||
@@ -4972,6 +4977,20 @@ impl AppState {
|
||||
self.pool_routing_control().close_removed(removed).await;
|
||||
}
|
||||
|
||||
pub async fn remove_plugin_connection_pools(&self, plugin_id: &str) {
|
||||
let removed = self.drain_plugin_connection_pools(plugin_id).await;
|
||||
self.pool_routing_control().close_removed(removed).await;
|
||||
}
|
||||
|
||||
pub async fn invoke_plugin_connection_action(
|
||||
&self,
|
||||
config: ConnectionConfig,
|
||||
action_id: &str,
|
||||
) -> Result<crate::plugins::PluginConnectionActionResult, String> {
|
||||
let (host, port) = self.connection_host_port(&config.id, &config).await?;
|
||||
self.plugin_host.invoke_connection_action(&config, action_id, &host, port).await
|
||||
}
|
||||
|
||||
async fn drain_connection_pools(&self, connection_id: &str) -> Vec<(String, PoolKind)> {
|
||||
let pool_prefix = format!("{connection_id}:");
|
||||
let keys_to_remove: Vec<String> = self
|
||||
@@ -5064,6 +5083,34 @@ impl AppState {
|
||||
removed
|
||||
}
|
||||
|
||||
async fn drain_plugin_connection_pools(&self, plugin_id: &str) -> Vec<(String, PoolKind)> {
|
||||
let keys_to_remove: Vec<String> = self
|
||||
.connections
|
||||
.read()
|
||||
.await
|
||||
.iter()
|
||||
.filter_map(|(key, pool)| match pool {
|
||||
PoolKind::PluginConnection(handle) if handle.plugin_id == plugin_id => Some(key.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
self.stop_keepalive_tasks(&keys_to_remove).await;
|
||||
{
|
||||
let mut activity = self.pool_activity.write().await;
|
||||
for key in &keys_to_remove {
|
||||
activity.remove(key);
|
||||
}
|
||||
}
|
||||
let mut conns = self.connections.write().await;
|
||||
let mut removed = Vec::with_capacity(keys_to_remove.len());
|
||||
for key in keys_to_remove {
|
||||
if let Some(pool) = conns.remove(&key) {
|
||||
removed.push((key, pool));
|
||||
}
|
||||
}
|
||||
removed
|
||||
}
|
||||
|
||||
async fn uses_forwarded_transport(&self, connection_id: &str) -> bool {
|
||||
let configs = self.configs.read().await;
|
||||
configs.get(connection_id).is_some_and(|config| config.has_effective_transport_layers())
|
||||
@@ -5602,7 +5649,6 @@ fn pool_key_for_session_role(
|
||||
|
||||
#[cfg(test)]
|
||||
fn clone_pool_kind(pool: &PoolKind) -> PoolKind {
|
||||
pool.clone()
|
||||
match pool {
|
||||
PoolKind::Mysql(p, mode) => PoolKind::Mysql(p.clone(), *mode),
|
||||
PoolKind::Postgres(p) => PoolKind::Postgres(p.clone()),
|
||||
@@ -5630,6 +5676,7 @@ fn clone_pool_kind(pool: &PoolKind) -> PoolKind {
|
||||
PoolKind::ExternalDriver { driver_id, config, session } => {
|
||||
PoolKind::ExternalDriver { driver_id: driver_id.clone(), config: config.clone(), session: session.clone() }
|
||||
}
|
||||
PoolKind::PluginConnection(handle) => PoolKind::PluginConnection(handle.clone()),
|
||||
PoolKind::MessageQueue => PoolKind::MessageQueue,
|
||||
PoolKind::Nacos => PoolKind::Nacos,
|
||||
PoolKind::Consul(client) => PoolKind::Consul(client.clone()),
|
||||
@@ -5639,18 +5686,6 @@ fn clone_pool_kind(pool: &PoolKind) -> PoolKind {
|
||||
}
|
||||
}
|
||||
|
||||
fn remove_mysql_pool_if_current(
|
||||
connections: &mut HashMap<String, PoolKind>,
|
||||
pool_key: &str,
|
||||
expected: &db::mysql::MySqlPool,
|
||||
) -> Option<PoolKind> {
|
||||
let is_current = matches!(
|
||||
connections.get(pool_key),
|
||||
Some(PoolKind::Mysql(current, _)) if expected.is_same_pool(current)
|
||||
);
|
||||
is_current.then(|| connections.remove(pool_key)).flatten()
|
||||
}
|
||||
|
||||
async fn close_pool_kind(pool: PoolKind) -> Result<(), String> {
|
||||
match pool {
|
||||
PoolKind::Mysql(p, _) => {
|
||||
@@ -5715,6 +5750,9 @@ async fn close_pool_kind(pool: PoolKind) -> Result<(), String> {
|
||||
PoolKind::ExternalDriver { session, .. } => {
|
||||
session.shutdown().await;
|
||||
}
|
||||
PoolKind::PluginConnection(handle) => {
|
||||
handle.disconnect().await?;
|
||||
}
|
||||
PoolKind::MessageQueue => {}
|
||||
PoolKind::Nacos => {}
|
||||
PoolKind::Consul(_) => {}
|
||||
|
||||
@@ -833,9 +833,6 @@ mod tests {
|
||||
CASSANDRA_TRUSTSTORE_PASSWORD_KEY, CONNECTION_STRING_KEY, INIT_SCRIPT_KEY, MAIN_PASSWORD_KEY,
|
||||
MQTT_AUTH_PASSWORD_KEY, MQ_AUTH_PASSWORD_KEY, MQ_AUTH_TOKEN_KEY, MQ_TOKEN_SIGNING_KEY,
|
||||
REDIS_SENTINEL_PASSWORD_KEY, SSH_PASSWORD_KEY,
|
||||
load_connections_from_file, save_connections_to_file, ConnectionSecretStore, CONNECTION_STRING_KEY,
|
||||
INIT_SCRIPT_KEY, MAIN_PASSWORD_KEY, MQTT_AUTH_PASSWORD_KEY, MQ_AUTH_PASSWORD_KEY, MQ_AUTH_TOKEN_KEY,
|
||||
MQ_TOKEN_SIGNING_KEY, REDIS_SENTINEL_PASSWORD_KEY, SSH_PASSWORD_KEY,
|
||||
};
|
||||
use crate::models::connection::{
|
||||
ConnectionConfig, DatabaseType, HttpTunnelConfig, SshTunnelConfig, TransportLayerConfig,
|
||||
|
||||
@@ -2246,6 +2246,7 @@ async fn do_execute_typed(
|
||||
.await
|
||||
.map(|result| truncate_result_with_max_rows(result, max_rows))
|
||||
}
|
||||
PoolKind::PluginConnection(_) => Err("SQL execution is not supported for plugin connections".to_string()),
|
||||
PoolKind::HBase(_) => Err("SQL execution is not supported for HBase connections".to_string()),
|
||||
PoolKind::DynamoDb(client) => {
|
||||
let client = client.clone();
|
||||
@@ -4022,7 +4023,8 @@ fn pool_kind_has_transactional_path(pool: &PoolKind) -> bool {
|
||||
| PoolKind::InfluxDb(_)
|
||||
| PoolKind::InfluxDb3(_)
|
||||
| PoolKind::VictoriaMetrics(_)
|
||||
| PoolKind::ExternalDriver { .. } => false,
|
||||
| PoolKind::ExternalDriver { .. }
|
||||
| PoolKind::PluginConnection(_) => false,
|
||||
#[cfg(feature = "mq-admin")]
|
||||
PoolKind::Mqtt(_) => false,
|
||||
}
|
||||
@@ -4276,34 +4278,6 @@ pub async fn execute_statements_in_transaction_on_pool_typed(
|
||||
// contract and rejects drivers that cannot provide it.
|
||||
// Clone the pool handle within the lock, then drop it before any async work.
|
||||
let path = { state.pool_handle(pool_key).await.as_ref().map(batch_transaction_path) };
|
||||
let path = {
|
||||
let conns = state.connections.read().await;
|
||||
conns.get(pool_key).map(|p| match p {
|
||||
PoolKind::Postgres(pg) => TxPath::Pg(pg.clone()),
|
||||
PoolKind::Mysql(mp, _mode) => TxPath::Mysql(mp.clone(), false),
|
||||
PoolKind::Sqlite(sq) => TxPath::Sqlite(sq.clone()),
|
||||
PoolKind::CloudflareD1(client) => TxPath::CloudflareD1(client.clone()),
|
||||
PoolKind::ClickHouse(_) | PoolKind::Rqlite(_) | PoolKind::Turso(_) | PoolKind::SqlServer(_) => {
|
||||
TxPath::Explicit
|
||||
}
|
||||
PoolKind::Agent(client) => TxPath::Agent(client.clone()),
|
||||
PoolKind::MessageQueue | PoolKind::Nacos | PoolKind::Consul(_) | PoolKind::HBase(_) => TxPath::None,
|
||||
#[cfg(feature = "mq-admin")]
|
||||
PoolKind::Mqtt(_) => TxPath::None,
|
||||
PoolKind::DuckDbWorker(_)
|
||||
| PoolKind::Redis(_)
|
||||
| PoolKind::MongoDb(_)
|
||||
| PoolKind::DynamoDb(_)
|
||||
| PoolKind::Elasticsearch(_)
|
||||
| PoolKind::Easysearch(_)
|
||||
| PoolKind::Meilisearch(_)
|
||||
| PoolKind::VectorDb(_)
|
||||
| PoolKind::InfluxDb(_)
|
||||
| PoolKind::InfluxDb3(_)
|
||||
| PoolKind::VictoriaMetrics(_)
|
||||
| PoolKind::ExternalDriver { .. } => TxPath::None,
|
||||
})
|
||||
};
|
||||
|
||||
let result = match path {
|
||||
Some(BatchTransactionPath::Pg(pool)) => {
|
||||
@@ -4420,7 +4394,8 @@ fn batch_transaction_path(pool: &PoolKind) -> BatchTransactionPath {
|
||||
| PoolKind::InfluxDb(_)
|
||||
| PoolKind::InfluxDb3(_)
|
||||
| PoolKind::VictoriaMetrics(_)
|
||||
| PoolKind::ExternalDriver { .. } => BatchTransactionPath::Unsupported,
|
||||
| PoolKind::ExternalDriver { .. }
|
||||
| PoolKind::PluginConnection(_) => BatchTransactionPath::Unsupported,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -817,9 +817,8 @@ pub fn profile_for(db_type: DatabaseType) -> DdlDialectProfile {
|
||||
}
|
||||
|
||||
InfluxDb => influxdb_profile(),
|
||||
DuckDb | Questdb | SapHana | Teradata | Snowflake | Trino | PrestoSql | Hive | Kyuubi | Impala | Argo
|
||||
| Spark | Db2 | Informix | Bigquery | Spanner | Kylin | Ignite | Ignite3 | Oscar | Tdengine | Iotdb
|
||||
| Databricks | Jdbc => conservative_ansi(db_type),
|
||||
|
||||
Plugin => conservative_ansi(db_type),
|
||||
|
||||
// Non-tabular / not applicable for relational CREATE TABLE
|
||||
Redis | MongoDb | DynamoDb | Elasticsearch | Easysearch | Meilisearch | Qdrant | Milvus | Weaviate
|
||||
|
||||
@@ -3279,8 +3279,8 @@ fn persist_connection_in_tx(tx: &rusqlite::Transaction<'_>, config: &ConnectionC
|
||||
persist_mq_auth_secrets_in_tx(tx, &config)?;
|
||||
persist_mq_token_signing_secret_in_tx(tx, &config)?;
|
||||
persist_nacos_auth_secrets_in_tx(tx, &config)?;
|
||||
persist_cassandra_tls_secrets_in_tx(tx, &config)
|
||||
persist_nacos_auth_secrets_in_tx(tx, &config)
|
||||
persist_cassandra_tls_secrets_in_tx(tx, &config)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn insert_connection_copy_next_to_source(entries: &mut Vec<serde_json::Value>, source_id: &str, copy_id: &str) -> bool {
|
||||
@@ -3742,8 +3742,6 @@ impl Storage {
|
||||
|| needs_mq_token_signing_rewrite
|
||||
|| needs_nacos_auth_rewrite
|
||||
|| needs_cassandra_tls_rewrite;
|
||||
let needs_external_secret_rewrite =
|
||||
needs_mq_auth_rewrite || needs_mq_token_signing_rewrite || needs_nacos_auth_rewrite;
|
||||
if needs_external_secret_rewrite {
|
||||
let mut sanitized = config.clone().canonicalized();
|
||||
scrub_mq_auth_secrets(&mut sanitized);
|
||||
@@ -5304,7 +5302,6 @@ mod tests {
|
||||
use crate::connection_secrets::{
|
||||
CASSANDRA_KEYSTORE_PASSWORD_KEY, CASSANDRA_TRUSTSTORE_PASSWORD_KEY, MQ_AUTH_PASSWORD_KEY, MQ_AUTH_TOKEN_KEY,
|
||||
MQ_TOKEN_SIGNING_KEY, NACOS_AUTH_PASSWORD_KEY,
|
||||
MQ_AUTH_PASSWORD_KEY, MQ_AUTH_TOKEN_KEY, MQ_TOKEN_SIGNING_KEY, NACOS_AUTH_PASSWORD_KEY,
|
||||
};
|
||||
use crate::history::{HistoryConnectionFilter, HistoryDatabaseFilter, HistoryEntry, HistorySearchRequest};
|
||||
use crate::models::connection::{
|
||||
|
||||
@@ -1811,6 +1811,33 @@ pub fn run() {
|
||||
commands::connection::save_sidebar_layout,
|
||||
commands::connection::load_sidebar_layout,
|
||||
commands::plugins::list_plugins,
|
||||
commands::plugins::list_plugin_trusted_keys,
|
||||
commands::plugins::save_plugin_trusted_key,
|
||||
commands::plugins::remove_plugin_trusted_key,
|
||||
commands::plugins::list_plugin_repositories,
|
||||
commands::plugins::save_plugin_repository,
|
||||
commands::plugins::remove_plugin_repository,
|
||||
commands::plugins::fetch_plugin_marketplace_catalogs,
|
||||
commands::plugins::install_marketplace_plugin,
|
||||
commands::plugins::install_plugin_package,
|
||||
commands::plugins::rollback_plugin,
|
||||
commands::plugins::uninstall_plugin,
|
||||
commands::plugins::activate_plugin,
|
||||
commands::plugins::list_active_plugins,
|
||||
commands::plugins::stop_plugin,
|
||||
commands::plugins::invoke_plugin,
|
||||
commands::plugins::invoke_plugin_connection_action,
|
||||
commands::plugins::notify_plugin,
|
||||
commands::plugins::send_plugin_binary,
|
||||
commands::plugins::list_plugin_filesystem_entries,
|
||||
commands::plugins::read_plugin_filesystem_file,
|
||||
commands::plugins::write_plugin_filesystem_file,
|
||||
commands::plugins::create_plugin_filesystem_directory,
|
||||
commands::plugins::delete_plugin_filesystem_entry,
|
||||
commands::plugins::rename_plugin_filesystem_entry,
|
||||
commands::plugins::read_plugin_ui_entry,
|
||||
commands::plugins::read_plugin_asset,
|
||||
commands::plugins::read_plugin_ui_asset,
|
||||
commands::plugins::list_jdbc_drivers,
|
||||
commands::plugins::list_jdbc_maven_bundles,
|
||||
commands::plugins::list_jdbc_local_bundles,
|
||||
|
||||
Reference in New Issue
Block a user