fix(plugins): wire plugin commands and restore build

This commit is contained in:
t8y2
2026-09-13 01:10:08 +08:00
parent d60ade2ac5
commit a97726e8e9
7 changed files with 90 additions and 60 deletions
-3
View File
@@ -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,
};
+54 -16
View File
@@ -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,
+5 -30
View File
@@ -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
+2 -5
View File
@@ -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::{
+27
View File
@@ -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,