feat(tdengine): replace Java agent with native Rust driver
This commit is contained in:
32 files changed
+5943
-1777
No files matched your search
@@ -0,0 +1,274 @@
|
||||
use anyhow::{bail, Context, Result};
|
||||
use percent_encoding::percent_decode_str;
|
||||
use std::net::IpAddr;
|
||||
use url::Url;
|
||||
|
||||
use crate::model::ConnectParams;
|
||||
|
||||
const DEFAULT_HOST: &str = "127.0.0.1";
|
||||
const DEFAULT_PORT: u16 = 6041;
|
||||
const DEFAULT_USER: &str = "root";
|
||||
const DEFAULT_PASSWORD: &str = "taosdata";
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct BuiltDsn {
|
||||
pub value: String,
|
||||
pub database: String,
|
||||
}
|
||||
|
||||
pub fn build_dsn(params: &ConnectParams) -> Result<BuiltDsn> {
|
||||
if !params.client_cert_path.trim().is_empty() || !params.client_key_path.trim().is_empty() {
|
||||
bail!("TDengine Rust WebSocket connector does not support client certificate authentication");
|
||||
}
|
||||
|
||||
let mut url = if params.connection_string.trim().is_empty() {
|
||||
build_from_fields(params)?
|
||||
} else {
|
||||
normalize_connection_string(params.connection_string.trim(), params.ssl)?
|
||||
};
|
||||
|
||||
apply_connection_fields(&mut url, params)?;
|
||||
merge_query_params(&mut url, ¶ms.url_params);
|
||||
if (params.ssl || !params.ca_cert_path.trim().is_empty()) && url.scheme() == "ws" {
|
||||
url.set_scheme("wss").map_err(|_| anyhow::anyhow!("failed to enable TLS in TDengine connection URL"))?;
|
||||
}
|
||||
if !params.ca_cert_path.trim().is_empty() {
|
||||
set_query_param(&mut url, "tls_mode", "verify_identity");
|
||||
set_query_param(&mut url, "tls_ca", params.ca_cert_path.trim());
|
||||
}
|
||||
let database = url
|
||||
.path_segments()
|
||||
.and_then(|mut segments| segments.find(|segment| !segment.is_empty()))
|
||||
.map(|segment| percent_decode_str(segment).decode_utf8_lossy().into_owned())
|
||||
.unwrap_or_default();
|
||||
Ok(BuiltDsn { value: url.into(), database })
|
||||
}
|
||||
|
||||
fn build_from_fields(params: &ConnectParams) -> Result<Url> {
|
||||
let scheme = if params.ssl { "wss" } else { "ws" };
|
||||
let host = if params.host.trim().is_empty() { DEFAULT_HOST } else { params.host.trim() };
|
||||
let port = if params.port == 0 { DEFAULT_PORT } else { params.port };
|
||||
let username = if params.username.is_empty() { DEFAULT_USER } else { ¶ms.username };
|
||||
let password = if params.password.is_empty() { DEFAULT_PASSWORD } else { ¶ms.password };
|
||||
let mut url = Url::parse(&format!("{scheme}://{DEFAULT_HOST}:{port}/"))?;
|
||||
if let Ok(address) = host.parse::<IpAddr>() {
|
||||
url.set_ip_host(address).map_err(|_| anyhow::anyhow!("invalid TDengine host"))?;
|
||||
} else {
|
||||
url.set_host(Some(host)).map_err(|_| anyhow::anyhow!("invalid TDengine host"))?;
|
||||
}
|
||||
url.set_username(username).map_err(|_| anyhow::anyhow!("invalid TDengine username"))?;
|
||||
url.set_password(Some(password)).map_err(|_| anyhow::anyhow!("invalid TDengine password"))?;
|
||||
if !params.database.trim().is_empty() {
|
||||
url.set_path(&format!("/{}", params.database.trim()));
|
||||
}
|
||||
Ok(url)
|
||||
}
|
||||
|
||||
fn apply_connection_fields(url: &mut Url, params: &ConnectParams) -> Result<()> {
|
||||
if url.username().is_empty() {
|
||||
let username = if params.username.is_empty() { DEFAULT_USER } else { ¶ms.username };
|
||||
url.set_username(username).map_err(|_| anyhow::anyhow!("invalid TDengine username"))?;
|
||||
}
|
||||
if url.password().is_none() {
|
||||
let password = if params.password.is_empty() { DEFAULT_PASSWORD } else { ¶ms.password };
|
||||
url.set_password(Some(password)).map_err(|_| anyhow::anyhow!("invalid TDengine password"))?;
|
||||
}
|
||||
if url.path().trim_matches('/').is_empty() && !params.database.trim().is_empty() {
|
||||
url.set_path(&format!("/{}", params.database.trim()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn normalize_connection_string(raw: &str, ssl: bool) -> Result<Url> {
|
||||
let trimmed = raw.trim();
|
||||
let normalized = if let Some(value) = strip_prefix_ignore_ascii_case(trimmed, "jdbc:TAOS-WS://") {
|
||||
format!("{}://{value}", if ssl { "wss" } else { "ws" })
|
||||
} else if let Some(value) = strip_prefix_ignore_ascii_case(trimmed, "jdbc:TAOS-RS://") {
|
||||
format!("{}://{value}", if ssl { "wss" } else { "ws" })
|
||||
} else if let Some(value) = strip_prefix_ignore_ascii_case(trimmed, "tdengine://") {
|
||||
format!("{}://{value}", if ssl { "wss" } else { "ws" })
|
||||
} else if let Some(value) = strip_prefix_ignore_ascii_case(trimmed, "taosws://") {
|
||||
format!("{}://{value}", if ssl { "wss" } else { "ws" })
|
||||
} else if let Some(value) = strip_prefix_ignore_ascii_case(trimmed, "taoswss://") {
|
||||
format!("wss://{value}")
|
||||
} else {
|
||||
trimmed.to_string()
|
||||
};
|
||||
let mut url = Url::parse(&normalized).with_context(|| "invalid TDengine connection string")?;
|
||||
if !matches!(url.scheme(), "ws" | "wss" | "http" | "https") {
|
||||
bail!("TDengine native agent supports only WebSocket connection strings");
|
||||
}
|
||||
if matches!(url.scheme(), "http") {
|
||||
url.set_scheme("ws").map_err(|_| anyhow::anyhow!("invalid TDengine HTTP connection string"))?;
|
||||
} else if matches!(url.scheme(), "https") {
|
||||
url.set_scheme("wss").map_err(|_| anyhow::anyhow!("invalid TDengine HTTPS connection string"))?;
|
||||
}
|
||||
remove_control_params(&mut url);
|
||||
Ok(url)
|
||||
}
|
||||
|
||||
fn merge_query_params(url: &mut Url, raw: &str) {
|
||||
let raw = raw.trim().trim_start_matches('?');
|
||||
if raw.is_empty() {
|
||||
return;
|
||||
}
|
||||
let additions = url::form_urlencoded::parse(raw.as_bytes())
|
||||
.filter(|(key, _)| !is_control_param(key))
|
||||
.map(|(key, value)| (key.into_owned(), value.into_owned()))
|
||||
.collect::<Vec<_>>();
|
||||
let mut query = url.query_pairs_mut();
|
||||
for (key, value) in additions {
|
||||
query.append_pair(&key, &value);
|
||||
}
|
||||
}
|
||||
|
||||
fn remove_control_params(url: &mut Url) {
|
||||
let kept = url
|
||||
.query_pairs()
|
||||
.filter(|(key, _)| !is_control_param(key))
|
||||
.map(|(key, value)| (key.into_owned(), value.into_owned()))
|
||||
.collect::<Vec<_>>();
|
||||
url.set_query(None);
|
||||
if kept.is_empty() {
|
||||
return;
|
||||
}
|
||||
let mut query = url.query_pairs_mut();
|
||||
for (key, value) in kept {
|
||||
query.append_pair(&key, &value);
|
||||
}
|
||||
}
|
||||
|
||||
fn set_query_param(url: &mut Url, key: &str, value: &str) {
|
||||
let mut pairs = url
|
||||
.query_pairs()
|
||||
.filter(|(existing, _)| !existing.eq_ignore_ascii_case(key))
|
||||
.map(|(existing, value)| (existing.into_owned(), value.into_owned()))
|
||||
.collect::<Vec<_>>();
|
||||
pairs.push((key.to_string(), value.to_string()));
|
||||
url.set_query(None);
|
||||
let mut query = url.query_pairs_mut();
|
||||
for (key, value) in pairs {
|
||||
query.append_pair(&key, &value);
|
||||
}
|
||||
}
|
||||
|
||||
fn is_control_param(key: &str) -> bool {
|
||||
key.eq_ignore_ascii_case("transport") || key.eq_ignore_ascii_case("dbx.transport")
|
||||
}
|
||||
|
||||
fn strip_prefix_ignore_ascii_case<'a>(value: &'a str, prefix: &str) -> Option<&'a str> {
|
||||
value.get(..prefix.len()).filter(|head| head.eq_ignore_ascii_case(prefix))?;
|
||||
value.get(prefix.len()..)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn builds_ws_dsn_from_connection_fields() {
|
||||
let dsn = build_dsn(&ConnectParams {
|
||||
host: "db.example.com".into(),
|
||||
port: 6041,
|
||||
database: "power_data".into(),
|
||||
username: "root".into(),
|
||||
password: "secret".into(),
|
||||
url_params: "timezone=UTC&dbx.transport=rest".into(),
|
||||
..Default::default()
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(dsn.value, "ws://root:secret@db.example.com:6041/power_data?timezone=UTC");
|
||||
assert_eq!(dsn.database, "power_data");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn accepts_legacy_jdbc_urls_without_transport_controls() {
|
||||
let dsn = build_dsn(&ConnectParams {
|
||||
connection_string: "jdbc:TAOS-WS://127.0.0.1:6041/db?transport=ws&timezone=UTC".into(),
|
||||
..Default::default()
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(dsn.value, "ws://root:taosdata@127.0.0.1:6041/db?timezone=UTC");
|
||||
assert_eq!(dsn.database, "db");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn fills_connection_string_credentials_and_database_from_fields() {
|
||||
let dsn = build_dsn(&ConnectParams {
|
||||
username: "reader".into(),
|
||||
password: "secret".into(),
|
||||
database: "metrics".into(),
|
||||
connection_string: "jdbc:TAOS-WS://td.example.com:6041?transport=ws".into(),
|
||||
..Default::default()
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(dsn.value, "ws://reader:secret@td.example.com:6041/metrics");
|
||||
assert_eq!(dsn.database, "metrics");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn preserves_explicit_connection_string_credentials_and_database() {
|
||||
let dsn = build_dsn(&ConnectParams {
|
||||
username: "ignored".into(),
|
||||
password: "ignored".into(),
|
||||
database: "ignored".into(),
|
||||
connection_string: "ws://url-user:url-pass@td.example.com:6041/url-db".into(),
|
||||
..Default::default()
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(dsn.value, "ws://url-user:url-pass@td.example.com:6041/url-db");
|
||||
assert_eq!(dsn.database, "url-db");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn passes_ca_certificate_path_to_the_websocket_connector() {
|
||||
let dsn =
|
||||
build_dsn(&ConnectParams { ssl: true, ca_cert_path: "/tmp/tdengine-ca.pem".into(), ..Default::default() })
|
||||
.unwrap();
|
||||
let url = Url::parse(&dsn.value).unwrap();
|
||||
assert_eq!(url.scheme(), "wss");
|
||||
let query = url.query_pairs().collect::<std::collections::HashMap<_, _>>();
|
||||
assert_eq!(query.get("tls_mode").map(|value| value.as_ref()), Some("verify_identity"));
|
||||
assert_eq!(query.get("tls_ca").map(|value| value.as_ref()), Some("/tmp/tdengine-ca.pem"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_mutual_tls_paths_instead_of_ignoring_them() {
|
||||
let error =
|
||||
build_dsn(&ConnectParams { client_cert_path: "/tmp/client.pem".into(), ..Default::default() }).unwrap_err();
|
||||
assert!(error.to_string().contains("client certificate"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn ca_certificate_enables_tls_and_replaces_conflicting_query_values() {
|
||||
let dsn = build_dsn(&ConnectParams {
|
||||
connection_string: "ws://td.example.com:6041/db?tls_mode=verify_ca&tls_ca=old.pem".into(),
|
||||
ca_cert_path: "/tmp/new-ca.pem".into(),
|
||||
..Default::default()
|
||||
})
|
||||
.unwrap();
|
||||
let url = Url::parse(&dsn.value).unwrap();
|
||||
assert_eq!(url.scheme(), "wss");
|
||||
let tls_mode = url.query_pairs().filter(|(key, _)| key == "tls_mode").collect::<Vec<_>>();
|
||||
let tls_ca = url.query_pairs().filter(|(key, _)| key == "tls_ca").collect::<Vec<_>>();
|
||||
assert_eq!(tls_mode.as_slice(), &[("tls_mode".into(), "verify_identity".into())]);
|
||||
assert_eq!(tls_ca.as_slice(), &[("tls_ca".into(), "/tmp/new-ca.pem".into())]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn supports_ipv6_hosts_and_decodes_connection_database() {
|
||||
let fields =
|
||||
build_dsn(&ConnectParams { host: "::1".into(), database: "power data".into(), ..Default::default() })
|
||||
.unwrap();
|
||||
assert_eq!(fields.value, "ws://root:taosdata@[::1]:6041/power%20data");
|
||||
assert_eq!(fields.database, "power data");
|
||||
|
||||
let connection = build_dsn(&ConnectParams {
|
||||
connection_string: "taoswss://td.example.com:6041/power%20data".into(),
|
||||
..Default::default()
|
||||
})
|
||||
.unwrap();
|
||||
assert_eq!(connection.value, "wss://root:taosdata@td.example.com:6041/power%20data");
|
||||
assert_eq!(connection.database, "power data");
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,7 @@
|
||||
mod config;
|
||||
mod driver;
|
||||
mod model;
|
||||
mod runtime;
|
||||
mod value;
|
||||
|
||||
pub use runtime::run;
|
||||
@@ -0,0 +1,7 @@
|
||||
#[tokio::main]
|
||||
async fn main() {
|
||||
if let Err(error) = dbx_tdengine_driver::run().await {
|
||||
eprintln!("TDengine driver failed: {error:#}");
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
@@ -1,980 +0,0 @@
|
||||
package com.dbx.agent.tdengine;
|
||||
|
||||
import com.dbx.agent.AbstractJdbcAgent;
|
||||
import com.dbx.agent.ColumnInfo;
|
||||
import com.dbx.agent.ConnectParams;
|
||||
import com.dbx.agent.DatabaseInfo;
|
||||
import com.dbx.agent.ExecuteQueryOptions;
|
||||
import com.dbx.agent.ForeignKeyInfo;
|
||||
import com.dbx.agent.IndexInfo;
|
||||
import com.dbx.agent.JdbcExecutor;
|
||||
import com.dbx.agent.MultiSessionJsonRpcServer;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.ObjectInfo;
|
||||
import com.dbx.agent.ObjectSource;
|
||||
import com.dbx.agent.QueryPageOptions;
|
||||
import com.dbx.agent.QueryPageResult;
|
||||
import com.dbx.agent.QueryResult;
|
||||
import com.dbx.agent.TableInfo;
|
||||
import com.dbx.agent.TriggerInfo;
|
||||
import com.taosdata.jdbc.TSDBDriver;
|
||||
import com.taosdata.jdbc.rs.RestfulConnection;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.sql.Connection;
|
||||
import java.sql.DriverManager;
|
||||
import java.sql.ResultSet;
|
||||
import java.sql.SQLException;
|
||||
import java.sql.Timestamp;
|
||||
import java.sql.Types;
|
||||
import java.time.format.DateTimeFormatter;
|
||||
import java.time.format.DateTimeFormatterBuilder;
|
||||
import java.time.temporal.ChronoField;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Collections;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.List;
|
||||
import java.util.Locale;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.Callable;
|
||||
import java.util.concurrent.ExecutionException;
|
||||
import java.util.concurrent.ExecutorService;
|
||||
import java.util.concurrent.Executors;
|
||||
import java.util.concurrent.Future;
|
||||
import java.util.concurrent.ThreadFactory;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.concurrent.TimeoutException;
|
||||
import java.util.regex.Matcher;
|
||||
import java.util.regex.Pattern;
|
||||
|
||||
public final class TDengineAgent extends AbstractJdbcAgent {
|
||||
private static final long TABLE_CACHE_TTL_MILLIS = 10_000L;
|
||||
private static final DateTimeFormatter TDENGINE_TIMESTAMP_FORMAT =
|
||||
new DateTimeFormatterBuilder()
|
||||
.appendPattern("yyyy-MM-dd HH:mm:ss")
|
||||
.appendFraction(ChronoField.NANO_OF_SECOND, 3, 9, true)
|
||||
.toFormatter();
|
||||
private static final Pattern NUMERIC_PRECISION_PATTERN =
|
||||
Pattern.compile("(?i)^(decimal|numeric)\\((\\d+)(?:,\\s*\\d+)?\\)");
|
||||
private static final Pattern NUMERIC_SCALE_PATTERN =
|
||||
Pattern.compile("(?i)^(decimal|numeric)\\(\\d+,\\s*(\\d+)\\)");
|
||||
private static final Pattern CHARACTER_LENGTH_PATTERN =
|
||||
Pattern.compile("(?i)^(binary|nchar|varchar|varbinary)\\((\\d+)\\)");
|
||||
private static final Pattern COMPOSITE_KEY_PATTERN =
|
||||
Pattern.compile("(?i)\\bCOMPOSITE\\s+KEY\\b");
|
||||
private static final Pattern TDENGINE_VERSION_PATTERN =
|
||||
Pattern.compile("^(\\d+)\\.(\\d+)\\.(\\d+)(?:\\.|$)");
|
||||
|
||||
private final Object tableCacheLock = new Object();
|
||||
private String tableCacheSchema = "";
|
||||
private long tableCacheTimeMillis;
|
||||
private List<TableInfo> tableCache = Collections.emptyList();
|
||||
private Connection restfulStateConnection;
|
||||
private String restfulOriginalCatalog;
|
||||
private String restfulOriginalDatabase;
|
||||
private boolean restInformationSchemaPagingEnabled;
|
||||
|
||||
@Override
|
||||
protected String driverClass() {
|
||||
return TDengineTransport.WEBSOCKET.driverClass();
|
||||
}
|
||||
|
||||
@Override
|
||||
protected String buildJdbcUrl(ConnectParams params) {
|
||||
String connectionString = params.getConnection_string() == null ? "" : params.getConnection_string().trim();
|
||||
if (!connectionString.isBlank()) {
|
||||
return TDengineJdbcUrl.sanitizeConnectionString(connectionString);
|
||||
}
|
||||
TDengineJdbcUrl.TransportPreference preference = TDengineJdbcUrl.transportPreference(params.getUrl_params());
|
||||
TDengineTransport transport = preference == TDengineJdbcUrl.TransportPreference.REST
|
||||
? TDengineTransport.REST
|
||||
: TDengineTransport.WEBSOCKET;
|
||||
return TDengineJdbcUrl.from(params, transport);
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void loadDriver(ConnectParams params) {
|
||||
}
|
||||
|
||||
@Override
|
||||
protected Connection openConnection(ConnectParams params) throws Exception {
|
||||
return TDengineConnectionFactory.open(params);
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void afterConnect(ConnectParams params, Connection connection) {
|
||||
clearTableCache();
|
||||
clearRestfulState();
|
||||
restInformationSchemaPagingEnabled = supportsRestInformationSchemaPaging(connection);
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void beforePooledConnectionReturn(Connection connection) throws Exception {
|
||||
if (restfulStateConnection != connection) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
restoreRestfulConnectionState(connection, restfulOriginalCatalog, restfulOriginalDatabase);
|
||||
} finally {
|
||||
clearRestfulState();
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void afterDisconnect() {
|
||||
clearTableCache();
|
||||
clearRestfulState();
|
||||
restInformationSchemaPagingEnabled = false;
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<DatabaseInfo> listDatabases() {
|
||||
return unchecked(() -> {
|
||||
List<DatabaseInfo> result = new ArrayList<>();
|
||||
try (java.sql.Statement stmt = requireConnected().createStatement();
|
||||
ResultSet rs = stmt.executeQuery("SHOW DATABASES")) {
|
||||
while (rs.next()) {
|
||||
String name = rs.getString(1);
|
||||
if (!isSystemDatabase(name)) {
|
||||
result.add(new DatabaseInfo(name));
|
||||
}
|
||||
}
|
||||
}
|
||||
return result;
|
||||
});
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<String> listSchemas() {
|
||||
return Collections.emptyList();
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<TableInfo> listTables(String schema) {
|
||||
return listTables(schema, MetadataListConstraints.NONE);
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<TableInfo> listTables(String schema, MetadataListConstraints constraints) {
|
||||
MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints);
|
||||
if (!normalized.tableTypeAllowed("TABLE")) {
|
||||
return Collections.emptyList();
|
||||
}
|
||||
if (normalized.hasLimit() && !normalized.hasFilter()) {
|
||||
if (restInformationSchemaPagingEnabled) {
|
||||
try {
|
||||
return listTablesPageFromRestInformationSchema(schema, normalized);
|
||||
} catch (RuntimeException ignored) {
|
||||
}
|
||||
}
|
||||
return listTablesPageFromShow(schema, normalized);
|
||||
}
|
||||
return normalized.filterTables(listTablesFromShow(schema));
|
||||
}
|
||||
|
||||
private List<TableInfo> listTablesPageFromRestInformationSchema(
|
||||
String schema,
|
||||
MetadataListConstraints constraints
|
||||
) {
|
||||
return unchecked(() -> {
|
||||
int limit = constraints.getLimit();
|
||||
int offset = constraints.getOffset() == null ? 0 : constraints.getOffset();
|
||||
List<TableInfo> result = new ArrayList<>(limit);
|
||||
|
||||
TablePageScan stablePage = queryTablesPage(
|
||||
"SHOW " + quoteQualifiedPrefix(schema) + "STABLES",
|
||||
"STABLE",
|
||||
false,
|
||||
offset,
|
||||
limit
|
||||
);
|
||||
result.addAll(stablePage.tables());
|
||||
if (result.size() >= limit) {
|
||||
return result;
|
||||
}
|
||||
|
||||
int tableOffset = Math.max(0, offset - stablePage.scannedRows());
|
||||
result.addAll(queryInformationSchemaTablePage(
|
||||
schema,
|
||||
tableOffset,
|
||||
limit - result.size()
|
||||
));
|
||||
return result;
|
||||
});
|
||||
}
|
||||
|
||||
private List<TableInfo> listTablesPageFromShow(String schema, MetadataListConstraints constraints) {
|
||||
return unchecked(() -> {
|
||||
int limit = constraints.getLimit();
|
||||
int offset = constraints.getOffset() == null ? 0 : constraints.getOffset();
|
||||
List<TableInfo> result = new ArrayList<>(limit);
|
||||
|
||||
TablePageScan stablePage = queryTablesPage(
|
||||
"SHOW " + quoteQualifiedPrefix(schema) + "STABLES",
|
||||
"STABLE",
|
||||
false,
|
||||
offset,
|
||||
limit
|
||||
);
|
||||
result.addAll(stablePage.tables());
|
||||
if (result.size() >= limit) {
|
||||
return result;
|
||||
}
|
||||
|
||||
int tableOffset = Math.max(0, offset - stablePage.scannedRows());
|
||||
result.addAll(queryTablesPage(
|
||||
"SHOW " + quoteQualifiedPrefix(schema) + "TABLES",
|
||||
"TABLE",
|
||||
true,
|
||||
tableOffset,
|
||||
limit - result.size()
|
||||
).tables());
|
||||
return result;
|
||||
});
|
||||
}
|
||||
|
||||
private List<TableInfo> listTablesFromShow(String schema) {
|
||||
List<TableInfo> cached = cachedTables(schema);
|
||||
if (cached != null) {
|
||||
return cached;
|
||||
}
|
||||
return unchecked(() -> {
|
||||
List<TableInfo> result = new ArrayList<>();
|
||||
// Connector/J 3.6.3 ignores Statement#setMaxRows. Filtered and
|
||||
// unbounded callers still need the complete result, so cache that
|
||||
// scan briefly instead of repeating it for adjacent requests.
|
||||
result.addAll(queryTables("SHOW " + quoteQualifiedPrefix(schema) + "STABLES", "STABLE", false));
|
||||
result.addAll(queryTables("SHOW " + quoteQualifiedPrefix(schema) + "TABLES", "TABLE", true));
|
||||
|
||||
Map<String, TableInfo> distinct = new LinkedHashMap<>();
|
||||
for (TableInfo table : result) {
|
||||
distinct.putIfAbsent(table.getTable_type() + ":" + table.getName(), table);
|
||||
}
|
||||
List<TableInfo> sorted = new ArrayList<>(distinct.values());
|
||||
sortTablesForHierarchy(sorted);
|
||||
cacheTables(schema, sorted);
|
||||
return copyTables(sorted);
|
||||
});
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<ObjectInfo> listObjects(String schema) {
|
||||
List<ObjectInfo> result = new ArrayList<>();
|
||||
for (TableInfo table : listTables(schema)) {
|
||||
result.add(new ObjectInfo(table.getName(), table.getTable_type(), schema, table.getComment()));
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<ColumnInfo> getColumns(String schema, String table) {
|
||||
return unchecked(() -> {
|
||||
try (java.sql.Statement stmt = requireConnected().createStatement();
|
||||
ResultSet rs = stmt.executeQuery("DESCRIBE " + qualifiedName(schema, table))) {
|
||||
return readDescribeColumns(rs);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
@Override
|
||||
public ObjectSource getObjectSource(String schema, String name, String objectType) {
|
||||
String source = getCreateSql(schema, name, objectType);
|
||||
if (source.isBlank()) {
|
||||
source = getTableDdl(schema, name);
|
||||
}
|
||||
return new ObjectSource(name, objectType, schema, source);
|
||||
}
|
||||
|
||||
@Override
|
||||
public String getTableDdl(String schema, String table) {
|
||||
String source = getCreateSql(schema, table, "STABLE");
|
||||
if (source.isBlank()) {
|
||||
source = getCreateSql(schema, table, "TABLE");
|
||||
}
|
||||
if (source.isBlank()) {
|
||||
return super.getTableDdl(schema, table);
|
||||
}
|
||||
return source;
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<IndexInfo> listIndexes(String schema, String table) {
|
||||
return Collections.emptyList();
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<ForeignKeyInfo> listForeignKeys(String schema, String table) {
|
||||
return Collections.emptyList();
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<TriggerInfo> listTriggers(String schema, String table) {
|
||||
return Collections.emptyList();
|
||||
}
|
||||
|
||||
@Override
|
||||
public QueryResult executeQuery(String sql, String schema, ExecuteQueryOptions options) {
|
||||
QueryResult result = JdbcExecutor.current().execute(
|
||||
requireConnected(),
|
||||
sql,
|
||||
prepareExecutionSchema(schema),
|
||||
this::setSchemaSQL,
|
||||
options.getMaxRows(),
|
||||
options.getFetchSize(),
|
||||
options.getTimeoutSecs(),
|
||||
this::resultValue
|
||||
);
|
||||
if (mayChangeMetadata(sql)) {
|
||||
clearTableCache();
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
@Override
|
||||
public QueryPageResult executeQueryPage(String sql, String schema, QueryPageOptions options) {
|
||||
return JdbcExecutor.current().executePage(
|
||||
requireConnected(),
|
||||
sql,
|
||||
prepareExecutionSchema(schema),
|
||||
this::setSchemaSQL,
|
||||
options,
|
||||
this::resultValue
|
||||
);
|
||||
}
|
||||
|
||||
@Override
|
||||
public QueryPageResult startTableRead(String sql, String schema, QueryPageOptions options) {
|
||||
return JdbcExecutor.current().startTableRead(
|
||||
requireConnected(),
|
||||
sql,
|
||||
prepareExecutionSchema(schema),
|
||||
this::setSchemaSQL,
|
||||
options,
|
||||
this::resultValue
|
||||
);
|
||||
}
|
||||
|
||||
@Override
|
||||
public String setSchemaSQL(String schema) {
|
||||
return "USE " + quoteIdentifier(schema);
|
||||
}
|
||||
|
||||
@Override
|
||||
public QueryResult executeTransaction(List<String> statements, String schema) {
|
||||
QueryResult result = super.executeTransaction(statements, prepareExecutionSchema(schema));
|
||||
if (statements.stream().anyMatch(TDengineAgent::mayChangeMetadata)) {
|
||||
clearTableCache();
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
@Override
|
||||
public QueryResult executeBatch(List<String> statements, String schema) {
|
||||
QueryResult result = super.executeBatch(statements, prepareExecutionSchema(schema));
|
||||
if (statements.stream().anyMatch(TDengineAgent::mayChangeMetadata)) {
|
||||
clearTableCache();
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
private String prepareExecutionSchema(String schema) {
|
||||
return unchecked(() -> {
|
||||
Connection connection = requireConnected();
|
||||
if (schema != null
|
||||
&& !schema.trim().isEmpty()
|
||||
&& unwrapConnection(connection, RestfulConnection.class) != null
|
||||
&& restfulStateConnection != connection) {
|
||||
restfulStateConnection = connection;
|
||||
restfulOriginalCatalog = connection.getCatalog();
|
||||
restfulOriginalDatabase = connection.getClientInfo(TSDBDriver.PROPERTY_KEY_DBNAME);
|
||||
}
|
||||
return prepareExecutionSchema(connection, schema);
|
||||
});
|
||||
}
|
||||
|
||||
static String prepareExecutionSchema(Connection connection, String schema) throws SQLException {
|
||||
if (unwrapConnection(connection, RestfulConnection.class) == null || schema == null || schema.trim().isEmpty()) {
|
||||
return schema;
|
||||
}
|
||||
|
||||
String database = schema.trim();
|
||||
// Connector/J 3.6.3 misparses quoted USE statements and puts the whole SQL in /rest/sql/<db>.
|
||||
connection.setCatalog(database);
|
||||
connection.setClientInfo(TSDBDriver.PROPERTY_KEY_DBNAME, database);
|
||||
return null;
|
||||
}
|
||||
|
||||
static void restoreRestfulConnectionState(Connection connection, String catalog, String database) throws SQLException {
|
||||
connection.setCatalog(catalog);
|
||||
if (database == null) {
|
||||
connection.getClientInfo().remove(TSDBDriver.PROPERTY_KEY_DBNAME);
|
||||
} else {
|
||||
connection.setClientInfo(TSDBDriver.PROPERTY_KEY_DBNAME, database);
|
||||
}
|
||||
}
|
||||
|
||||
private List<TableInfo> queryTables(String sql, String tableType, boolean includesStableName) throws Exception {
|
||||
List<TableInfo> result = new ArrayList<>();
|
||||
try (java.sql.Statement stmt = requireConnected().createStatement()) {
|
||||
try (ResultSet rs = stmt.executeQuery(sql)) {
|
||||
while (rs.next()) {
|
||||
result.add(readShowTable(rs, tableType, includesStableName));
|
||||
}
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
private TablePageScan queryTablesPage(
|
||||
String sql,
|
||||
String tableType,
|
||||
boolean includesStableName,
|
||||
int offset,
|
||||
int limit
|
||||
) throws Exception {
|
||||
List<TableInfo> result = new ArrayList<>(limit);
|
||||
int scannedRows = 0;
|
||||
try (java.sql.Statement stmt = requireConnected().createStatement();
|
||||
ResultSet rs = stmt.executeQuery(sql)) {
|
||||
while (result.size() < limit && rs.next()) {
|
||||
scannedRows += 1;
|
||||
if (scannedRows <= offset) {
|
||||
continue;
|
||||
}
|
||||
result.add(readShowTable(rs, tableType, includesStableName));
|
||||
}
|
||||
}
|
||||
return new TablePageScan(result, scannedRows);
|
||||
}
|
||||
|
||||
private List<TableInfo> queryInformationSchemaTablePage(String database, int offset, int limit) throws Exception {
|
||||
StringBuilder sql = new StringBuilder(
|
||||
"SELECT table_name, stable_name FROM information_schema.ins_tables WHERE db_name = ? LIMIT "
|
||||
).append(limit);
|
||||
if (offset > 0) {
|
||||
sql.append(" OFFSET ").append(offset);
|
||||
}
|
||||
|
||||
List<TableInfo> result = new ArrayList<>(limit);
|
||||
try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement(sql.toString())) {
|
||||
stmt.setString(1, database);
|
||||
try (ResultSet rs = stmt.executeQuery()) {
|
||||
while (rs.next()) {
|
||||
String parentName = optionalString(rs, 2);
|
||||
if (parentName != null && parentName.trim().isEmpty()) {
|
||||
parentName = null;
|
||||
}
|
||||
result.add(new TableInfo(rs.getString(1), "TABLE", null, null, parentName));
|
||||
}
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
private static TableInfo readShowTable(ResultSet rs, String tableType, boolean includesStableName) throws Exception {
|
||||
// SHOW TABLES returns the owning STABLE as its fourth column. It
|
||||
// is absent for ordinary tables and older servers, where this
|
||||
// best-effort read simply leaves the table at the root level.
|
||||
String parentName = includesStableName ? optionalString(rs, 4) : null;
|
||||
if (parentName != null && parentName.trim().isEmpty()) {
|
||||
parentName = null;
|
||||
}
|
||||
return new TableInfo(rs.getString(1), tableType, null, null, parentName);
|
||||
}
|
||||
|
||||
private record TablePageScan(List<TableInfo> tables, int scannedRows) {
|
||||
}
|
||||
|
||||
private static boolean supportsRestInformationSchemaPaging(Connection connection) {
|
||||
if (unwrapConnection(connection, RestfulConnection.class) == null) {
|
||||
return false;
|
||||
}
|
||||
try {
|
||||
return supportsRestInformationSchemaPagingVersion(connection.getMetaData().getDatabaseProductVersion());
|
||||
} catch (Exception ignored) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
static boolean supportsRestInformationSchemaPagingVersion(String version) {
|
||||
Matcher matcher = TDENGINE_VERSION_PATTERN.matcher(version == null ? "" : version.trim());
|
||||
if (!matcher.find()) {
|
||||
return false;
|
||||
}
|
||||
int major = Integer.parseInt(matcher.group(1));
|
||||
int minor = Integer.parseInt(matcher.group(2));
|
||||
int patch = Integer.parseInt(matcher.group(3));
|
||||
return major > 3 || (major == 3 && (minor > 3 || (minor == 3 && patch >= 8)));
|
||||
}
|
||||
|
||||
private List<TableInfo> cachedTables(String schema) {
|
||||
String normalizedSchema = normalizedSchema(schema);
|
||||
synchronized (tableCacheLock) {
|
||||
if (cacheFresh(tableCacheTimeMillis) && tableCacheSchema.equals(normalizedSchema)) {
|
||||
return copyTables(tableCache);
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
private void cacheTables(String schema, List<TableInfo> tables) {
|
||||
synchronized (tableCacheLock) {
|
||||
tableCacheSchema = normalizedSchema(schema);
|
||||
tableCache = copyTables(tables);
|
||||
tableCacheTimeMillis = System.currentTimeMillis();
|
||||
}
|
||||
}
|
||||
|
||||
private void clearTableCache() {
|
||||
synchronized (tableCacheLock) {
|
||||
tableCacheSchema = "";
|
||||
tableCacheTimeMillis = 0L;
|
||||
tableCache = Collections.emptyList();
|
||||
}
|
||||
}
|
||||
|
||||
private void clearRestfulState() {
|
||||
restfulStateConnection = null;
|
||||
restfulOriginalCatalog = null;
|
||||
restfulOriginalDatabase = null;
|
||||
}
|
||||
|
||||
private static boolean cacheFresh(long cachedAtMillis) {
|
||||
return cachedAtMillis > 0L && System.currentTimeMillis() - cachedAtMillis <= TABLE_CACHE_TTL_MILLIS;
|
||||
}
|
||||
|
||||
private static String normalizedSchema(String schema) {
|
||||
return schema == null ? "" : schema.trim();
|
||||
}
|
||||
|
||||
private static List<TableInfo> copyTables(List<TableInfo> tables) {
|
||||
List<TableInfo> copies = new ArrayList<>(tables.size());
|
||||
for (TableInfo table : tables) {
|
||||
copies.add(new TableInfo(
|
||||
table.getName(),
|
||||
table.getTable_type(),
|
||||
table.getComment(),
|
||||
table.getParent_schema(),
|
||||
table.getParent_name()
|
||||
));
|
||||
}
|
||||
return copies;
|
||||
}
|
||||
|
||||
private static boolean mayChangeMetadata(String sql) {
|
||||
String normalized = sql == null ? "" : sql.trim().toLowerCase(Locale.ROOT);
|
||||
return normalized.startsWith("create ")
|
||||
|| normalized.startsWith("drop ")
|
||||
|| normalized.startsWith("alter ")
|
||||
|| normalized.startsWith("rename ")
|
||||
|| normalized.startsWith("truncate ");
|
||||
}
|
||||
|
||||
static void sortTablesForHierarchy(List<TableInfo> tables) {
|
||||
tables.sort((left, right) -> {
|
||||
String leftGroup = hierarchyGroupName(left);
|
||||
String rightGroup = hierarchyGroupName(right);
|
||||
int groupCompared = leftGroup.compareTo(rightGroup);
|
||||
if (groupCompared != 0) {
|
||||
return groupCompared;
|
||||
}
|
||||
|
||||
boolean leftIsParent = left.getParent_name() == null;
|
||||
boolean rightIsParent = right.getParent_name() == null;
|
||||
if (leftIsParent != rightIsParent) {
|
||||
return leftIsParent ? -1 : 1;
|
||||
}
|
||||
return left.getName().toLowerCase(Locale.ROOT).compareTo(right.getName().toLowerCase(Locale.ROOT));
|
||||
});
|
||||
}
|
||||
|
||||
private static String hierarchyGroupName(TableInfo table) {
|
||||
String parentName = table.getParent_name();
|
||||
String groupName = parentName == null || parentName.trim().isEmpty() ? table.getName() : parentName;
|
||||
return groupName.toLowerCase(Locale.ROOT);
|
||||
}
|
||||
|
||||
static List<ColumnInfo> readDescribeColumns(ResultSet rs) throws Exception {
|
||||
List<ColumnInfo> result = new ArrayList<>();
|
||||
int ordinal = 0;
|
||||
while (rs.next()) {
|
||||
ordinal += 1;
|
||||
String name = rs.getString(1);
|
||||
String dataType = coalesce(rs.getString(2));
|
||||
String note = optionalString(rs, 4);
|
||||
boolean isTag = note != null && note.toUpperCase(Locale.ROOT).contains("TAG");
|
||||
boolean isPrimaryKey = !isTag && (ordinal == 1 || (note != null && COMPOSITE_KEY_PATTERN.matcher(note).find()));
|
||||
result.add(new ColumnInfo(
|
||||
name,
|
||||
dataType,
|
||||
!isPrimaryKey,
|
||||
null,
|
||||
isPrimaryKey,
|
||||
note,
|
||||
isTag ? "TAG" : null,
|
||||
parseNumericPrecision(dataType),
|
||||
parseNumericScale(dataType),
|
||||
parseCharacterMaximumLength(dataType)
|
||||
));
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
private String getCreateSql(String schema, String name, String objectType) {
|
||||
String showType = switch (objectType.toUpperCase(Locale.ROOT)) {
|
||||
case "STABLE", "SUPER TABLE", "SUPERTABLE" -> "STABLE";
|
||||
case "TABLE", "BASE TABLE", "CHILD TABLE" -> "TABLE";
|
||||
default -> null;
|
||||
};
|
||||
if (showType == null) {
|
||||
return "";
|
||||
}
|
||||
|
||||
try (java.sql.Statement stmt = requireConnected().createStatement();
|
||||
ResultSet rs = stmt.executeQuery("SHOW CREATE " + showType + " " + qualifiedName(schema, name))) {
|
||||
if (!rs.next()) {
|
||||
return "";
|
||||
}
|
||||
int columnCount = rs.getMetaData().getColumnCount();
|
||||
for (int i = 2; i <= columnCount; i++) {
|
||||
String value = rs.getString(i);
|
||||
if (value != null && value.toUpperCase(Locale.ROOT).contains("CREATE")) {
|
||||
return value;
|
||||
}
|
||||
}
|
||||
String value = rs.getString(columnCount);
|
||||
return value == null ? "" : value;
|
||||
} catch (Exception e) {
|
||||
return "";
|
||||
}
|
||||
}
|
||||
|
||||
@Override
|
||||
protected Object resultValue(ResultSet rs, int index, int sqlType) {
|
||||
return unchecked(() -> {
|
||||
Object value = switch (sqlType) {
|
||||
case Types.BIGINT -> rs.getLong(index);
|
||||
case Types.INTEGER, Types.SMALLINT, Types.TINYINT -> rs.getInt(index);
|
||||
case Types.FLOAT, Types.REAL -> rs.getFloat(index);
|
||||
case Types.DOUBLE -> rs.getDouble(index);
|
||||
case Types.DECIMAL, Types.NUMERIC -> rs.getBigDecimal(index);
|
||||
case Types.BOOLEAN, Types.BIT -> rs.getBoolean(index);
|
||||
default -> rs.getObject(index);
|
||||
};
|
||||
return rs.wasNull() ? null : decodeTdengineValue(value);
|
||||
});
|
||||
}
|
||||
|
||||
public static Object decodeTdengineValue(Object value) {
|
||||
if (value instanceof byte[] bytes) {
|
||||
return new String(bytes, StandardCharsets.UTF_8);
|
||||
}
|
||||
if (value instanceof Timestamp timestamp) {
|
||||
return TDENGINE_TIMESTAMP_FORMAT.format(timestamp.toLocalDateTime());
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
private static String quoteQualifiedPrefix(String schema) {
|
||||
String trimmed = schema.trim();
|
||||
return trimmed.isEmpty() ? "" : quoteIdentifier(trimmed) + ".";
|
||||
}
|
||||
|
||||
private static String qualifiedName(String schema, String name) {
|
||||
String table = quoteIdentifier(name);
|
||||
String trimmed = schema.trim();
|
||||
return trimmed.isEmpty() ? table : quoteIdentifier(trimmed) + "." + table;
|
||||
}
|
||||
|
||||
private static String quoteIdentifier(String identifier) {
|
||||
return "`" + identifier.replace("`", "``") + "`";
|
||||
}
|
||||
|
||||
private static boolean isSystemDatabase(String name) {
|
||||
if (name == null) {
|
||||
return false;
|
||||
}
|
||||
String normalized = name.trim().toLowerCase(Locale.ROOT);
|
||||
return "information_schema".equals(normalized) || "performance_schema".equals(normalized);
|
||||
}
|
||||
|
||||
private static Integer parseNumericPrecision(String dataType) {
|
||||
return parseIntGroup(NUMERIC_PRECISION_PATTERN, dataType, 2);
|
||||
}
|
||||
|
||||
private static Integer parseNumericScale(String dataType) {
|
||||
return parseIntGroup(NUMERIC_SCALE_PATTERN, dataType, 2);
|
||||
}
|
||||
|
||||
private static Integer parseCharacterMaximumLength(String dataType) {
|
||||
return parseIntGroup(CHARACTER_LENGTH_PATTERN, dataType, 2);
|
||||
}
|
||||
|
||||
private static Integer parseIntGroup(Pattern pattern, String value, int group) {
|
||||
Matcher matcher = pattern.matcher(value);
|
||||
if (!matcher.find()) {
|
||||
return null;
|
||||
}
|
||||
try {
|
||||
return Integer.valueOf(matcher.group(group));
|
||||
} catch (NumberFormatException e) {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
private static String optionalString(ResultSet rs, int index) {
|
||||
try {
|
||||
return rs.getString(index);
|
||||
} catch (Exception e) {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
private static String coalesce(String value) {
|
||||
return value == null ? "" : value;
|
||||
}
|
||||
|
||||
public static void main(String[] args) {
|
||||
new MultiSessionJsonRpcServer(TDengineAgent::new).run();
|
||||
}
|
||||
}
|
||||
|
||||
final class TDengineJdbcUrl {
|
||||
enum TransportPreference {
|
||||
AUTO,
|
||||
WEBSOCKET,
|
||||
REST
|
||||
}
|
||||
|
||||
private static final String PARAM_TRANSPORT = "transport";
|
||||
private static final String PARAM_DBX_TRANSPORT = "dbx.transport";
|
||||
|
||||
private TDengineJdbcUrl() {
|
||||
}
|
||||
|
||||
static String from(ConnectParams params) {
|
||||
TransportPreference preference = transportPreference(params.getUrl_params());
|
||||
if (preference == TransportPreference.REST) {
|
||||
return from(params, TDengineTransport.REST);
|
||||
}
|
||||
String host = params.getHost().isBlank() ? "localhost" : params.getHost();
|
||||
int port = params.getPort() > 0 ? params.getPort() : 6041;
|
||||
String database = params.getDatabase().trim();
|
||||
String path = database.isBlank() ? "/" : "/" + database;
|
||||
String query = normalizeQueryString(params.getUrl_params());
|
||||
String suffix = query.isBlank() ? "" : "?" + query;
|
||||
return "jdbc:TAOS-WS://" + host + ":" + port + path + suffix;
|
||||
}
|
||||
|
||||
static String from(ConnectParams params, TDengineTransport transport) {
|
||||
String host = params.getHost().isBlank() ? "localhost" : params.getHost();
|
||||
int port = params.getPort() > 0 ? params.getPort() : 6041;
|
||||
String database = params.getDatabase().trim();
|
||||
String path = database.isBlank() ? "/" : "/" + database;
|
||||
String query = normalizeQueryString(params.getUrl_params());
|
||||
String suffix = query.isBlank() ? "" : "?" + query;
|
||||
return transport.urlPrefix() + host + ":" + port + path + suffix;
|
||||
}
|
||||
|
||||
static String sanitizeConnectionString(String connectionString) {
|
||||
if (connectionString == null || connectionString.isBlank()) {
|
||||
return "";
|
||||
}
|
||||
int queryIndex = connectionString.indexOf('?');
|
||||
if (queryIndex < 0) {
|
||||
return connectionString;
|
||||
}
|
||||
String base = connectionString.substring(0, queryIndex);
|
||||
String rawQuery = connectionString.substring(queryIndex + 1);
|
||||
int fragmentIndex = rawQuery.indexOf('#');
|
||||
String fragment = "";
|
||||
if (fragmentIndex >= 0) {
|
||||
fragment = rawQuery.substring(fragmentIndex);
|
||||
rawQuery = rawQuery.substring(0, fragmentIndex);
|
||||
}
|
||||
String query = normalizeQueryString(rawQuery);
|
||||
if (query.isBlank()) {
|
||||
return base + fragment;
|
||||
}
|
||||
return base + "?" + query + fragment;
|
||||
}
|
||||
|
||||
static TransportPreference transportPreference(String rawQuery) {
|
||||
String query = rawQuery == null ? "" : rawQuery.trim();
|
||||
if (query.startsWith("?")) {
|
||||
query = query.substring(1);
|
||||
}
|
||||
if (query.isBlank()) {
|
||||
return TransportPreference.AUTO;
|
||||
}
|
||||
String[] parts = query.split("&");
|
||||
for (String part : parts) {
|
||||
if (part == null || part.isBlank()) {
|
||||
continue;
|
||||
}
|
||||
int index = part.indexOf('=');
|
||||
String key = index < 0 ? part.trim() : part.substring(0, index).trim();
|
||||
if (!PARAM_TRANSPORT.equalsIgnoreCase(key) && !PARAM_DBX_TRANSPORT.equalsIgnoreCase(key)) {
|
||||
continue;
|
||||
}
|
||||
String value = index < 0 ? "" : part.substring(index + 1).trim();
|
||||
if ("rest".equalsIgnoreCase(value) || "rs".equalsIgnoreCase(value)) {
|
||||
return TransportPreference.REST;
|
||||
}
|
||||
if ("ws".equalsIgnoreCase(value) || "websocket".equalsIgnoreCase(value)) {
|
||||
return TransportPreference.WEBSOCKET;
|
||||
}
|
||||
}
|
||||
return TransportPreference.AUTO;
|
||||
}
|
||||
|
||||
private static String normalizeQueryString(String rawQuery) {
|
||||
String query = rawQuery == null ? "" : rawQuery.trim();
|
||||
if (query.startsWith("?")) {
|
||||
query = query.substring(1);
|
||||
}
|
||||
if (query.isBlank()) {
|
||||
return "";
|
||||
}
|
||||
String[] parts = query.split("&");
|
||||
List<String> kept = new ArrayList<>();
|
||||
for (String part : parts) {
|
||||
if (part == null || part.isBlank()) {
|
||||
continue;
|
||||
}
|
||||
int index = part.indexOf('=');
|
||||
String key = index < 0 ? part.trim() : part.substring(0, index).trim();
|
||||
if (PARAM_TRANSPORT.equalsIgnoreCase(key) || PARAM_DBX_TRANSPORT.equalsIgnoreCase(key)) {
|
||||
continue;
|
||||
}
|
||||
kept.add(part);
|
||||
}
|
||||
return String.join("&", kept);
|
||||
}
|
||||
}
|
||||
|
||||
enum TDengineTransport {
|
||||
WEBSOCKET("jdbc:TAOS-WS://", "com.taosdata.jdbc.ws.WebSocketDriver"),
|
||||
REST("jdbc:TAOS-RS://", "com.taosdata.jdbc.rs.RestfulDriver");
|
||||
|
||||
private final String urlPrefix;
|
||||
private final String driverClass;
|
||||
|
||||
TDengineTransport(String urlPrefix, String driverClass) {
|
||||
this.urlPrefix = urlPrefix;
|
||||
this.driverClass = driverClass;
|
||||
}
|
||||
|
||||
String urlPrefix() {
|
||||
return urlPrefix;
|
||||
}
|
||||
|
||||
String driverClass() {
|
||||
return driverClass;
|
||||
}
|
||||
}
|
||||
|
||||
final class TDengineConnectionFactory {
|
||||
private static final int AUTO_TRANSPORT_ATTEMPT_TIMEOUT_SECONDS = 8;
|
||||
|
||||
private TDengineConnectionFactory() {
|
||||
}
|
||||
|
||||
static Connection open(ConnectParams params) throws Exception {
|
||||
String explicitConnectionString = params.getConnection_string() == null ? "" : params.getConnection_string().trim();
|
||||
if (!explicitConnectionString.isBlank()) {
|
||||
String sanitizedConnectionString = TDengineJdbcUrl.sanitizeConnectionString(explicitConnectionString);
|
||||
Class.forName(driverClassFromConnectionString(sanitizedConnectionString));
|
||||
return DriverManager.getConnection(sanitizedConnectionString, params.getUsername(), params.getPassword());
|
||||
}
|
||||
|
||||
TDengineJdbcUrl.TransportPreference preference = TDengineJdbcUrl.transportPreference(params.getUrl_params());
|
||||
TDengineTransport[] candidates = transportsFor(preference);
|
||||
Exception firstError = null;
|
||||
for (TDengineTransport transport : candidates) {
|
||||
try {
|
||||
Connection connection = preference == TDengineJdbcUrl.TransportPreference.AUTO
|
||||
? openWithAttemptTimeout(params, transport)
|
||||
: openDirect(params, transport);
|
||||
return connection;
|
||||
} catch (Exception e) {
|
||||
if (firstError == null) {
|
||||
firstError = e;
|
||||
} else {
|
||||
firstError.addSuppressed(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
if (firstError != null) {
|
||||
throw new RuntimeException("Failed to connect TDengine using transports: " + candidateNames(candidates), firstError);
|
||||
}
|
||||
throw new IllegalStateException("No TDengine transport available.");
|
||||
}
|
||||
|
||||
private static Connection openDirect(ConnectParams params, TDengineTransport transport) throws Exception {
|
||||
Class.forName(transport.driverClass());
|
||||
String url = TDengineJdbcUrl.from(params, transport);
|
||||
return DriverManager.getConnection(url, params.getUsername(), params.getPassword());
|
||||
}
|
||||
|
||||
private static Connection openWithAttemptTimeout(ConnectParams params, TDengineTransport transport) throws Exception {
|
||||
String threadName = "tdengine-connect-" + transport.name().toLowerCase(Locale.ROOT);
|
||||
ThreadFactory factory = runnable -> {
|
||||
Thread thread = new Thread(runnable, threadName);
|
||||
thread.setDaemon(true);
|
||||
return thread;
|
||||
};
|
||||
ExecutorService executor = Executors.newSingleThreadExecutor(factory);
|
||||
try {
|
||||
Future<Connection> future = executor.submit(new Callable<Connection>() {
|
||||
@Override
|
||||
public Connection call() throws Exception {
|
||||
return openDirect(params, transport);
|
||||
}
|
||||
});
|
||||
return future.get(AUTO_TRANSPORT_ATTEMPT_TIMEOUT_SECONDS, TimeUnit.SECONDS);
|
||||
} catch (TimeoutException e) {
|
||||
throw new RuntimeException(
|
||||
"TDengine " + transport.name().toLowerCase(Locale.ROOT)
|
||||
+ " transport connect timed out after " + AUTO_TRANSPORT_ATTEMPT_TIMEOUT_SECONDS + "s",
|
||||
e
|
||||
);
|
||||
} catch (InterruptedException e) {
|
||||
Thread.currentThread().interrupt();
|
||||
throw new RuntimeException("Interrupted while connecting TDengine via " + transport.name().toLowerCase(Locale.ROOT), e);
|
||||
} catch (ExecutionException e) {
|
||||
Throwable cause = e.getCause();
|
||||
if (cause instanceof Exception exception) {
|
||||
throw exception;
|
||||
}
|
||||
throw new RuntimeException(cause);
|
||||
} finally {
|
||||
executor.shutdownNow();
|
||||
}
|
||||
}
|
||||
|
||||
private static TDengineTransport[] transportsFor(TDengineJdbcUrl.TransportPreference preference) {
|
||||
return switch (preference) {
|
||||
case REST -> new TDengineTransport[] {TDengineTransport.REST};
|
||||
case WEBSOCKET -> new TDengineTransport[] {TDengineTransport.WEBSOCKET};
|
||||
case AUTO -> new TDengineTransport[] {TDengineTransport.WEBSOCKET, TDengineTransport.REST};
|
||||
};
|
||||
}
|
||||
|
||||
private static String candidateNames(TDengineTransport[] candidates) {
|
||||
List<String> names = new ArrayList<>();
|
||||
for (TDengineTransport candidate : candidates) {
|
||||
names.add(candidate.name().toLowerCase(Locale.ROOT));
|
||||
}
|
||||
return String.join(" -> ", names);
|
||||
}
|
||||
|
||||
private static String driverClassFromConnectionString(String connectionString) {
|
||||
if (connectionString.regionMatches(true, 0, TDengineTransport.REST.urlPrefix(), 0, TDengineTransport.REST.urlPrefix().length())) {
|
||||
return TDengineTransport.REST.driverClass();
|
||||
}
|
||||
if (connectionString.regionMatches(true, 0, TDengineTransport.WEBSOCKET.urlPrefix(), 0, TDengineTransport.WEBSOCKET.urlPrefix().length())) {
|
||||
return TDengineTransport.WEBSOCKET.driverClass();
|
||||
}
|
||||
return TDengineTransport.WEBSOCKET.driverClass();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,279 @@
|
||||
use serde::{Deserialize, Deserializer, Serialize};
|
||||
use serde_json::Value;
|
||||
|
||||
pub const PROTOCOL_VERSION: u32 = 2;
|
||||
pub const LEGACY_SESSION_ID: &str = "__legacy__";
|
||||
pub const DEFAULT_MAX_ROWS: usize = 10_000;
|
||||
pub const MAX_AGENT_SESSIONS: usize = 256;
|
||||
pub const MAX_CONCURRENT_REQUESTS: usize = 64;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct RpcRequest {
|
||||
#[serde(default)]
|
||||
pub id: Value,
|
||||
pub method: String,
|
||||
#[serde(default)]
|
||||
pub params: Value,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct RpcResponse {
|
||||
pub jsonrpc: &'static str,
|
||||
pub id: Value,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub result: Option<Value>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub error: Option<RpcError>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct RpcError {
|
||||
pub code: i32,
|
||||
pub message: String,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub data: Option<StructuredError>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct StructuredError {
|
||||
pub category: &'static str,
|
||||
pub retryable: bool,
|
||||
pub session_disposition: &'static str,
|
||||
pub stage: &'static str,
|
||||
pub contract_version: u32,
|
||||
pub operation_outcome: &'static str,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub agent_session_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub exception_class: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
pub struct ConnectParams {
|
||||
#[serde(default)]
|
||||
pub host: String,
|
||||
#[serde(default)]
|
||||
pub port: u16,
|
||||
#[serde(default)]
|
||||
pub database: String,
|
||||
#[serde(default)]
|
||||
pub username: String,
|
||||
#[serde(default)]
|
||||
pub password: String,
|
||||
#[serde(default)]
|
||||
pub url_params: String,
|
||||
#[serde(default)]
|
||||
pub connection_string: String,
|
||||
#[serde(default)]
|
||||
pub ssl: bool,
|
||||
#[serde(default)]
|
||||
pub ca_cert_path: String,
|
||||
#[serde(default)]
|
||||
pub client_cert_path: String,
|
||||
#[serde(default)]
|
||||
pub client_key_path: String,
|
||||
#[serde(default)]
|
||||
pub connect_timeout_secs: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
pub struct QueryOptions {
|
||||
#[serde(default)]
|
||||
pub sql: String,
|
||||
#[serde(default)]
|
||||
pub database: String,
|
||||
#[serde(default)]
|
||||
pub schema: String,
|
||||
#[serde(default, rename = "maxRows")]
|
||||
pub max_rows: usize,
|
||||
#[serde(default, rename = "fetchSize")]
|
||||
pub fetch_size: usize,
|
||||
#[serde(default, rename = "timeoutSecs")]
|
||||
pub timeout_secs: u64,
|
||||
#[serde(default, rename = "pageSize")]
|
||||
pub page_size: usize,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
pub struct MetadataListConstraints {
|
||||
#[serde(default)]
|
||||
pub filter: String,
|
||||
#[serde(default)]
|
||||
pub limit: usize,
|
||||
#[serde(default)]
|
||||
pub offset: usize,
|
||||
#[serde(default)]
|
||||
pub object_types: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
pub struct CompletionAssistantRequest {
|
||||
#[serde(default)]
|
||||
pub database: String,
|
||||
#[serde(default, deserialize_with = "deserialize_null_default")]
|
||||
pub schema: String,
|
||||
#[serde(default)]
|
||||
pub object_kinds: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub mask: String,
|
||||
#[serde(default)]
|
||||
pub case_sensitive: bool,
|
||||
#[serde(default)]
|
||||
pub global_search: bool,
|
||||
#[serde(default, deserialize_with = "deserialize_null_default")]
|
||||
pub max_results: usize,
|
||||
#[serde(default, deserialize_with = "deserialize_null_default")]
|
||||
pub parent_schema: String,
|
||||
#[serde(default, deserialize_with = "deserialize_null_default")]
|
||||
pub parent_name: String,
|
||||
#[serde(default, deserialize_with = "deserialize_null_default")]
|
||||
pub match_mode: String,
|
||||
}
|
||||
|
||||
fn deserialize_null_default<'de, D, T>(deserializer: D) -> Result<T, D::Error>
|
||||
where
|
||||
D: Deserializer<'de>,
|
||||
T: Deserialize<'de> + Default,
|
||||
{
|
||||
Ok(Option::<T>::deserialize(deserializer)?.unwrap_or_default())
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct HandshakeResult {
|
||||
#[serde(rename = "protocolVersion")]
|
||||
pub protocol_version: u32,
|
||||
#[serde(rename = "agentProtocolVersion")]
|
||||
pub agent_protocol_version: u32,
|
||||
pub capabilities: Vec<&'static str>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
|
||||
pub struct DatabaseInfo {
|
||||
pub name: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
|
||||
pub struct TableInfo {
|
||||
pub name: String,
|
||||
pub table_type: String,
|
||||
pub comment: Option<String>,
|
||||
pub parent_schema: Option<String>,
|
||||
pub parent_name: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
|
||||
pub struct ObjectInfo {
|
||||
pub name: String,
|
||||
pub object_type: String,
|
||||
pub schema: String,
|
||||
pub comment: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
|
||||
pub struct ColumnInfo {
|
||||
pub name: String,
|
||||
pub data_type: String,
|
||||
pub is_nullable: bool,
|
||||
pub column_default: Option<String>,
|
||||
pub is_primary_key: bool,
|
||||
pub extra: Option<String>,
|
||||
pub comment: Option<String>,
|
||||
pub numeric_precision: Option<u32>,
|
||||
pub numeric_scale: Option<u32>,
|
||||
pub character_maximum_length: Option<u32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
|
||||
pub struct ObjectSource {
|
||||
pub name: String,
|
||||
pub object_type: String,
|
||||
pub schema: String,
|
||||
pub source: String,
|
||||
pub editable: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
|
||||
pub struct CompletionAssistantCandidate {
|
||||
pub name: String,
|
||||
pub kind: String,
|
||||
pub database: Option<String>,
|
||||
pub schema: Option<String>,
|
||||
pub parent_schema: Option<String>,
|
||||
pub parent_name: Option<String>,
|
||||
pub comment: Option<String>,
|
||||
pub data_type: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
|
||||
pub struct CompletionAssistantResponse {
|
||||
pub candidates: Vec<CompletionAssistantCandidate>,
|
||||
pub incomplete: bool,
|
||||
pub fallback_used: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Serialize, PartialEq)]
|
||||
pub struct QueryResult {
|
||||
pub columns: Vec<String>,
|
||||
pub column_types: Vec<String>,
|
||||
pub rows: Vec<Vec<Value>>,
|
||||
pub affected_rows: i64,
|
||||
pub execution_time_ms: i64,
|
||||
pub truncated: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Serialize, PartialEq)]
|
||||
pub struct QueryPageResult {
|
||||
pub columns: Vec<String>,
|
||||
pub column_types: Vec<String>,
|
||||
pub rows: Vec<Vec<Value>>,
|
||||
pub affected_rows: i64,
|
||||
pub execution_time_ms: i64,
|
||||
pub truncated: bool,
|
||||
pub session_id: Option<String>,
|
||||
pub has_more: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct DatabaseConnectionInfo {
|
||||
pub product_name: Option<String>,
|
||||
pub product_version: Option<String>,
|
||||
pub current_database: Option<String>,
|
||||
pub driver_name: Option<String>,
|
||||
pub driver_version: Option<String>,
|
||||
pub unquoted_identifier_case: Option<&'static str>,
|
||||
pub quoted_identifier_case: Option<&'static str>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AgentConnectionInfo {
|
||||
pub identifier_quote: String,
|
||||
pub compatibility_mode: Option<String>,
|
||||
pub database_info: Option<DatabaseConnectionInfo>,
|
||||
}
|
||||
|
||||
pub const TDENGINE_DATA_TYPES: &[&str] = &[
|
||||
"TIMESTAMP",
|
||||
"BOOL",
|
||||
"TINYINT",
|
||||
"SMALLINT",
|
||||
"INT",
|
||||
"BIGINT",
|
||||
"TINYINT UNSIGNED",
|
||||
"SMALLINT UNSIGNED",
|
||||
"INT UNSIGNED",
|
||||
"BIGINT UNSIGNED",
|
||||
"FLOAT",
|
||||
"DOUBLE",
|
||||
"BINARY",
|
||||
"VARCHAR",
|
||||
"NCHAR",
|
||||
"JSON",
|
||||
"VARBINARY",
|
||||
"GEOMETRY",
|
||||
"DECIMAL",
|
||||
"BLOB",
|
||||
"MEDIUMBLOB",
|
||||
];
|
||||
@@ -0,0 +1,662 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::{Arc, Mutex as StdMutex};
|
||||
|
||||
use anyhow::{anyhow, bail, Context, Result};
|
||||
use serde::de::DeserializeOwned;
|
||||
use serde_json::{json, Value};
|
||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
|
||||
use tokio::sync::{mpsc, Mutex, OwnedSemaphorePermit, RwLock, Semaphore};
|
||||
use tokio::task::JoinSet;
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
use crate::driver::TdengineDriver;
|
||||
use crate::model::{
|
||||
ConnectParams, HandshakeResult, RpcError, RpcRequest, RpcResponse, StructuredError, LEGACY_SESSION_ID,
|
||||
MAX_AGENT_SESSIONS, MAX_CONCURRENT_REQUESTS, PROTOCOL_VERSION,
|
||||
};
|
||||
|
||||
struct RuntimeServer {
|
||||
sessions: RwLock<HashMap<String, Arc<AgentSession>>>,
|
||||
session_slots: Arc<Semaphore>,
|
||||
request_slots: Arc<Semaphore>,
|
||||
}
|
||||
|
||||
struct AgentSession {
|
||||
driver: Mutex<TdengineDriver>,
|
||||
active: StdMutex<Option<ActiveOperation>>,
|
||||
next_operation_id: AtomicU64,
|
||||
_slot: OwnedSemaphorePermit,
|
||||
}
|
||||
|
||||
struct ActiveOperation {
|
||||
id: u64,
|
||||
token: CancellationToken,
|
||||
}
|
||||
|
||||
struct OperationGuard<'a> {
|
||||
session: &'a AgentSession,
|
||||
id: u64,
|
||||
}
|
||||
|
||||
impl Drop for OperationGuard<'_> {
|
||||
fn drop(&mut self) {
|
||||
let mut active = self.session.active.lock().expect("active operation lock poisoned");
|
||||
if active.as_ref().is_some_and(|operation| operation.id == self.id) {
|
||||
*active = None;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AgentSession {
|
||||
fn begin_operation(&self) -> (CancellationToken, OperationGuard<'_>) {
|
||||
let id = self.next_operation_id.fetch_add(1, Ordering::Relaxed);
|
||||
let token = CancellationToken::new();
|
||||
*self.active.lock().expect("active operation lock poisoned") =
|
||||
Some(ActiveOperation { id, token: token.clone() });
|
||||
(token, OperationGuard { session: self, id })
|
||||
}
|
||||
|
||||
fn cancel(&self) {
|
||||
if let Some(operation) = self.active.lock().expect("active operation lock poisoned").as_ref() {
|
||||
operation.token.cancel();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run() -> Result<()> {
|
||||
let runtime = Arc::new(RuntimeServer::new());
|
||||
let (responses, mut response_rx) = mpsc::unbounded_channel::<RpcResponse>();
|
||||
let writer = tokio::spawn(async move {
|
||||
let mut stdout = tokio::io::stdout();
|
||||
stdout.write_all(b"{\"ready\":true}\n").await?;
|
||||
stdout.flush().await?;
|
||||
while let Some(response) = response_rx.recv().await {
|
||||
let line = serde_json::to_vec(&response)?;
|
||||
stdout.write_all(&line).await?;
|
||||
stdout.write_all(b"\n").await?;
|
||||
stdout.flush().await?;
|
||||
}
|
||||
Ok::<(), anyhow::Error>(())
|
||||
});
|
||||
|
||||
let mut lines = BufReader::new(tokio::io::stdin()).lines();
|
||||
let mut requests = JoinSet::new();
|
||||
while let Some(line) = lines.next_line().await? {
|
||||
while requests.try_join_next().is_some() {}
|
||||
let line = line.trim();
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let parsed = serde_json::from_str::<RpcRequest>(line);
|
||||
if parsed.as_ref().is_ok_and(|request| request.method == "shutdown") {
|
||||
let response = match parsed {
|
||||
Ok(request) => handle_request(runtime.clone(), request).await,
|
||||
Err(error) => error_response(Value::Null, "request", None, error.into()),
|
||||
};
|
||||
responses.send(response).map_err(|_| anyhow!("TDengine response writer stopped"))?;
|
||||
while requests.join_next().await.is_some() {}
|
||||
break;
|
||||
}
|
||||
let request_permit = if parsed.as_ref().is_ok_and(|request| is_capacity_exempt(&request.method)) {
|
||||
None
|
||||
} else {
|
||||
match runtime.request_slots.clone().try_acquire_owned() {
|
||||
Ok(permit) => Some(permit),
|
||||
Err(_) => {
|
||||
let response = match parsed {
|
||||
Ok(request) => error_response(
|
||||
if request.id.is_null() { json!(1) } else { request.id },
|
||||
&request.method,
|
||||
session_id(&request.params),
|
||||
anyhow!("agent request capacity is temporarily exhausted"),
|
||||
),
|
||||
Err(error) => error_response(Value::Null, "request", None, error.into()),
|
||||
};
|
||||
responses.send(response).map_err(|_| anyhow!("TDengine response writer stopped"))?;
|
||||
continue;
|
||||
}
|
||||
}
|
||||
};
|
||||
let runtime = runtime.clone();
|
||||
let responses = responses.clone();
|
||||
requests.spawn(async move {
|
||||
let _request_permit = request_permit;
|
||||
let response = match parsed {
|
||||
Ok(request) => handle_request(runtime, request).await,
|
||||
Err(error) => error_response(Value::Null, "request", None, error.into()),
|
||||
};
|
||||
let _ = responses.send(response);
|
||||
});
|
||||
}
|
||||
runtime.close_all_sessions().await;
|
||||
while requests.join_next().await.is_some() {}
|
||||
drop(responses);
|
||||
writer.await.context("TDengine response writer task failed")??;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_request(runtime: Arc<RuntimeServer>, request: RpcRequest) -> RpcResponse {
|
||||
let id = if request.id.is_null() { json!(1) } else { request.id.clone() };
|
||||
let session_id = session_id(&request.params);
|
||||
match runtime.dispatch(&request.method, request.params).await {
|
||||
Ok(result) => RpcResponse { jsonrpc: "2.0", id, result: Some(result), error: None },
|
||||
Err(error) => error_response(id, &request.method, session_id, error),
|
||||
}
|
||||
}
|
||||
|
||||
impl RuntimeServer {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
sessions: RwLock::new(HashMap::new()),
|
||||
session_slots: Arc::new(Semaphore::new(MAX_AGENT_SESSIONS)),
|
||||
request_slots: Arc::new(Semaphore::new(MAX_CONCURRENT_REQUESTS)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn dispatch(&self, method: &str, params: Value) -> Result<Value> {
|
||||
match method {
|
||||
"handshake" => serialize(HandshakeResult {
|
||||
protocol_version: PROTOCOL_VERSION,
|
||||
agent_protocol_version: PROTOCOL_VERSION,
|
||||
capabilities: vec![
|
||||
"connect",
|
||||
"test_connection",
|
||||
"metadata",
|
||||
"query",
|
||||
"paged_query",
|
||||
"transaction",
|
||||
"ddl",
|
||||
"multi_session",
|
||||
"structured_error_v1",
|
||||
],
|
||||
}),
|
||||
"open_session" => {
|
||||
let id = required_string(¶ms, "agentSessionId")?;
|
||||
let connect = decode::<ConnectParams>(¶ms)?;
|
||||
self.open_session(&id, connect).await?;
|
||||
Ok(json!({"ok": true}))
|
||||
}
|
||||
"close_session" => {
|
||||
self.close_session(&required_string(¶ms, "agentSessionId")?).await?;
|
||||
Ok(json!({"ok": true}))
|
||||
}
|
||||
"validate_session" => {
|
||||
let session = self.session(&required_string(¶ms, "agentSessionId")?).await?;
|
||||
let driver = session.driver.lock().await;
|
||||
let (token, _operation) = session.begin_operation();
|
||||
driver.validate_connection(&token).await?;
|
||||
Ok(json!({"ok": true}))
|
||||
}
|
||||
"cancel_session" => {
|
||||
self.session(&required_string(¶ms, "agentSessionId")?).await?.cancel();
|
||||
Ok(json!({"ok": true}))
|
||||
}
|
||||
"test_connection" => {
|
||||
let database_info = TdengineDriver::test_connection(decode::<ConnectParams>(¶ms)?).await?;
|
||||
Ok(json!({"ok": true, "databaseInfo": database_info}))
|
||||
}
|
||||
"connect" => {
|
||||
self.close_session(LEGACY_SESSION_ID).await?;
|
||||
self.open_session(LEGACY_SESSION_ID, decode::<ConnectParams>(¶ms)?).await?;
|
||||
Ok(json!({"ok": true}))
|
||||
}
|
||||
"disconnect" => {
|
||||
self.close_session(LEGACY_SESSION_ID).await?;
|
||||
Ok(json!({"ok": true}))
|
||||
}
|
||||
"shutdown" => {
|
||||
self.close_all_sessions().await;
|
||||
Ok(json!({"ok": true}))
|
||||
}
|
||||
_ => {
|
||||
if !is_driver_method(method) {
|
||||
bail!("unknown method: {method}");
|
||||
}
|
||||
let id = session_id(¶ms).unwrap_or_else(|| LEGACY_SESSION_ID.to_string());
|
||||
let session = self.session(&id).await?;
|
||||
let mut driver = session.driver.lock().await;
|
||||
let (token, _operation) = session.begin_operation();
|
||||
dispatch_driver(&mut driver, method, ¶ms, &token).await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn open_session(&self, id: &str, params: ConnectParams) -> Result<()> {
|
||||
if id.trim().is_empty() {
|
||||
bail!("agentSessionId is required");
|
||||
}
|
||||
{
|
||||
let sessions = self.sessions.read().await;
|
||||
if sessions.contains_key(id) {
|
||||
bail!("agent session already exists: {id}");
|
||||
}
|
||||
}
|
||||
let slot = self
|
||||
.session_slots
|
||||
.clone()
|
||||
.try_acquire_owned()
|
||||
.map_err(|_| anyhow!("agent session limit reached: {MAX_AGENT_SESSIONS}"))?;
|
||||
let mut driver = TdengineDriver::new();
|
||||
driver.connect(params).await?;
|
||||
let session = Arc::new(AgentSession {
|
||||
driver: Mutex::new(driver),
|
||||
active: StdMutex::new(None),
|
||||
next_operation_id: AtomicU64::new(1),
|
||||
_slot: slot,
|
||||
});
|
||||
let mut sessions = self.sessions.write().await;
|
||||
if sessions.contains_key(id) {
|
||||
bail!("agent session already exists: {id}");
|
||||
}
|
||||
sessions.insert(id.to_string(), session);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn session(&self, id: &str) -> Result<Arc<AgentSession>> {
|
||||
self.sessions.read().await.get(id).cloned().ok_or_else(|| anyhow!("agent session not found: {id}"))
|
||||
}
|
||||
|
||||
async fn close_session(&self, id: &str) -> Result<()> {
|
||||
let session = self.sessions.write().await.remove(id);
|
||||
let Some(session) = session else {
|
||||
return Ok(());
|
||||
};
|
||||
session.cancel();
|
||||
session.driver.lock().await.disconnect().await;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn close_all_sessions(&self) {
|
||||
let sessions = {
|
||||
let mut sessions = self.sessions.write().await;
|
||||
std::mem::take(&mut *sessions)
|
||||
};
|
||||
for (_, session) in sessions {
|
||||
session.cancel();
|
||||
session.driver.lock().await.disconnect().await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn dispatch_driver(
|
||||
driver: &mut TdengineDriver,
|
||||
method: &str,
|
||||
params: &Value,
|
||||
token: &CancellationToken,
|
||||
) -> Result<Value> {
|
||||
let scope = database_scope(params);
|
||||
match method {
|
||||
"validate_connection" => {
|
||||
driver.validate_connection(token).await?;
|
||||
Ok(json!({"ok": true}))
|
||||
}
|
||||
"connection_info" => serialize(driver.connection_info()?),
|
||||
"list_databases" => serialize(driver.list_databases(token).await?),
|
||||
"list_schemas" => serialize(driver.list_schemas()),
|
||||
"list_tables" => serialize(driver.list_tables(&scope, decode(params)?, token).await?),
|
||||
"get_table_comment" => serialize(driver.get_table_comment(&scope, &required_string(params, "table")?).await?),
|
||||
"list_objects" => serialize(driver.list_objects(&scope, decode(params)?, token).await?),
|
||||
"list_data_types" => serialize(driver.list_data_types()),
|
||||
"completion_assistant_search_v1" => {
|
||||
serialize(driver.completion_assistant_search(decode(params)?, token).await?)
|
||||
}
|
||||
"get_columns" => serialize(driver.get_columns(&scope, &required_string(params, "table")?, token).await?),
|
||||
"list_indexes" | "list_foreign_keys" | "list_triggers" | "list_constraints" | "list_partitions"
|
||||
| "list_subpartitions" => Ok(json!([])),
|
||||
"get_object_source" => serialize(
|
||||
driver
|
||||
.get_object_source(
|
||||
&scope,
|
||||
&required_string(params, "name")?,
|
||||
&required_string(params, "object_type")?,
|
||||
token,
|
||||
)
|
||||
.await?,
|
||||
),
|
||||
"get_table_ddl" => serialize(driver.get_table_ddl(&scope, &required_string(params, "table")?, token).await?),
|
||||
"get_explain_info" => Ok(json!({
|
||||
"plan": driver
|
||||
.get_explain_info(
|
||||
&required_string(params, "sql")?,
|
||||
token,
|
||||
optional_u64(params, "timeoutSecs").unwrap_or(0),
|
||||
)
|
||||
.await?,
|
||||
"has_actual_stats": false,
|
||||
})),
|
||||
"execute_query" => serialize(driver.execute_query(decode(params)?, token).await?),
|
||||
"execute_query_page" | "start_table_read" => {
|
||||
serialize(driver.execute_query_page(decode(params)?, token).await?)
|
||||
}
|
||||
"fetch_query_page" | "fetch_table_read_page" => serialize(
|
||||
driver
|
||||
.fetch_query_page(
|
||||
&required_string(params, "sessionId")?,
|
||||
optional_usize(params, "pageSize").unwrap_or(100),
|
||||
optional_u64(params, "timeoutSecs").unwrap_or(0),
|
||||
token,
|
||||
)
|
||||
.await?,
|
||||
),
|
||||
"close_query_session" | "close_table_read_session" => {
|
||||
serialize(driver.close_query_session(&required_string(params, "sessionId")?))
|
||||
}
|
||||
"execute_transaction" | "execute_batch" => serialize(
|
||||
driver
|
||||
.execute_statements(
|
||||
string_array(params, "statements")?,
|
||||
&scope,
|
||||
optional_u64(params, "timeoutSecs").unwrap_or(0),
|
||||
token,
|
||||
)
|
||||
.await?,
|
||||
),
|
||||
"disconnect" => {
|
||||
driver.disconnect().await;
|
||||
Ok(json!({"ok": true}))
|
||||
}
|
||||
"shutdown" => {
|
||||
driver.disconnect().await;
|
||||
Ok(json!({"ok": true}))
|
||||
}
|
||||
_ => bail!("unknown method: {method}"),
|
||||
}
|
||||
}
|
||||
|
||||
fn decode<T: DeserializeOwned>(params: &Value) -> Result<T> {
|
||||
serde_json::from_value(params.clone()).with_context(|| "invalid TDengine agent request parameters")
|
||||
}
|
||||
|
||||
fn serialize<T: serde::Serialize>(value: T) -> Result<Value> {
|
||||
serde_json::to_value(value).map_err(anyhow::Error::from)
|
||||
}
|
||||
|
||||
fn required_string(params: &Value, key: &str) -> Result<String> {
|
||||
params
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string)
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.ok_or_else(|| anyhow!("{key} is required"))
|
||||
}
|
||||
|
||||
fn optional_string(params: &Value, key: &str) -> String {
|
||||
params.get(key).and_then(Value::as_str).unwrap_or_default().to_string()
|
||||
}
|
||||
|
||||
fn optional_usize(params: &Value, key: &str) -> Option<usize> {
|
||||
params.get(key).and_then(Value::as_u64).and_then(|value| usize::try_from(value).ok())
|
||||
}
|
||||
|
||||
fn optional_u64(params: &Value, key: &str) -> Option<u64> {
|
||||
params.get(key).and_then(Value::as_u64)
|
||||
}
|
||||
|
||||
fn string_array(params: &Value, key: &str) -> Result<Vec<String>> {
|
||||
params
|
||||
.get(key)
|
||||
.cloned()
|
||||
.ok_or_else(|| anyhow!("{key} is required"))
|
||||
.and_then(|value| serde_json::from_value(value).map_err(anyhow::Error::from))
|
||||
}
|
||||
|
||||
fn database_scope(params: &Value) -> String {
|
||||
let schema = optional_string(params, "schema");
|
||||
if schema.trim().is_empty() {
|
||||
optional_string(params, "database")
|
||||
} else {
|
||||
schema
|
||||
}
|
||||
}
|
||||
|
||||
fn session_id(params: &Value) -> Option<String> {
|
||||
params.get("agentSessionId").and_then(Value::as_str).map(str::to_string)
|
||||
}
|
||||
|
||||
fn is_driver_method(method: &str) -> bool {
|
||||
matches!(
|
||||
method,
|
||||
"validate_connection"
|
||||
| "connection_info"
|
||||
| "list_databases"
|
||||
| "list_schemas"
|
||||
| "list_tables"
|
||||
| "get_table_comment"
|
||||
| "list_objects"
|
||||
| "list_data_types"
|
||||
| "completion_assistant_search_v1"
|
||||
| "get_columns"
|
||||
| "list_indexes"
|
||||
| "list_foreign_keys"
|
||||
| "list_triggers"
|
||||
| "list_constraints"
|
||||
| "list_partitions"
|
||||
| "list_subpartitions"
|
||||
| "get_object_source"
|
||||
| "get_table_ddl"
|
||||
| "get_explain_info"
|
||||
| "execute_query"
|
||||
| "execute_query_page"
|
||||
| "start_table_read"
|
||||
| "fetch_query_page"
|
||||
| "fetch_table_read_page"
|
||||
| "close_query_session"
|
||||
| "close_table_read_session"
|
||||
| "execute_transaction"
|
||||
| "execute_batch"
|
||||
)
|
||||
}
|
||||
|
||||
fn is_capacity_exempt(method: &str) -> bool {
|
||||
matches!(method, "handshake" | "cancel_session" | "close_session" | "disconnect")
|
||||
}
|
||||
|
||||
fn error_response(id: Value, method: &str, session_id: Option<String>, error: anyhow::Error) -> RpcResponse {
|
||||
let message = format!("{error:#}");
|
||||
let lower = message.to_ascii_lowercase();
|
||||
let stage = error_stage(method);
|
||||
let (category, retryable, session_disposition) = if lower.contains("cancel") {
|
||||
("canceled", false, "quarantine")
|
||||
} else if lower.contains("timed out") || lower.contains("timeout") {
|
||||
("timeout", false, "quarantine")
|
||||
} else if lower.contains("memory exhausted")
|
||||
|| lower.contains("resource exhausted")
|
||||
|| lower.contains("session limit reached")
|
||||
|| lower.contains("request capacity")
|
||||
{
|
||||
("resource", true, "keep")
|
||||
} else if lower.contains("agent session not found")
|
||||
|| lower.contains("query session not found")
|
||||
|| lower.contains("agent session already exists")
|
||||
|| lower.contains("unknown method")
|
||||
|| lower.contains("invalid tdengine agent request parameters")
|
||||
|| lower.contains(" is required")
|
||||
{
|
||||
("protocol", false, "keep")
|
||||
} else if stage == "connect" || stage == "validate" || lower.contains("connection") || lower.contains("websocket") {
|
||||
(
|
||||
"connection",
|
||||
stage == "connect" || stage == "validate",
|
||||
if stage == "connect" { "keep" } else { "quarantine" },
|
||||
)
|
||||
} else if matches!(stage, "execute" | "fetch") {
|
||||
("sql", false, "keep")
|
||||
} else {
|
||||
("protocol", false, "keep")
|
||||
};
|
||||
RpcResponse {
|
||||
jsonrpc: "2.0",
|
||||
id,
|
||||
result: None,
|
||||
error: Some(RpcError {
|
||||
code: -1,
|
||||
message,
|
||||
data: Some(StructuredError {
|
||||
category,
|
||||
retryable,
|
||||
session_disposition,
|
||||
stage,
|
||||
contract_version: 1,
|
||||
operation_outcome: if matches!(stage, "request" | "connect" | "validate") {
|
||||
"not_started"
|
||||
} else {
|
||||
"unknown"
|
||||
},
|
||||
agent_session_id: session_id,
|
||||
exception_class: Some("dbx_tdengine_driver::Error".into()),
|
||||
}),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn error_stage(method: &str) -> &'static str {
|
||||
match method {
|
||||
"request" | "handshake" => "request",
|
||||
"connect" | "open_session" | "test_connection" => "connect",
|
||||
"validate_connection" | "validate_session" => "validate",
|
||||
"cancel_session" => "cancel",
|
||||
"close_session" | "disconnect" | "close_query_session" | "close_table_read_session" | "shutdown" => "close",
|
||||
"fetch_query_page" | "fetch_table_read_page" => "fetch",
|
||||
_ => "execute",
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn structured_errors_keep_stage_and_outcome_consistent() {
|
||||
let response = error_response(json!(1), "open_session", Some("s1".into()), anyhow!("connection refused"));
|
||||
let data = response.error.unwrap().data.unwrap();
|
||||
assert_eq!(data.category, "connection");
|
||||
assert_eq!(data.stage, "connect");
|
||||
assert_eq!(data.operation_outcome, "not_started");
|
||||
assert_eq!(data.agent_session_id.as_deref(), Some("s1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn classifies_tdengine_execution_failures_as_sql_errors() {
|
||||
let response = error_response(json!(1), "execute_query", Some("s1".into()), anyhow!("invalid function"));
|
||||
let data = response.error.unwrap().data.unwrap();
|
||||
assert_eq!(data.category, "sql");
|
||||
assert_eq!(data.stage, "execute");
|
||||
assert_eq!(data.operation_outcome, "unknown");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_invalid_request_parameters_as_protocol_errors() {
|
||||
let response = error_response(
|
||||
json!(1),
|
||||
"execute_query",
|
||||
Some("s1".into()),
|
||||
anyhow!("invalid TDengine agent request parameters"),
|
||||
);
|
||||
assert_eq!(response.error.unwrap().data.unwrap().category, "protocol");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_missing_query_sessions_as_protocol_errors() {
|
||||
let response =
|
||||
error_response(json!(1), "fetch_query_page", Some("s1".into()), anyhow!("query session not found"));
|
||||
let data = response.error.unwrap().data.unwrap();
|
||||
assert_eq!(data.category, "protocol");
|
||||
assert_eq!(data.stage, "fetch");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quarantines_connections_lost_during_execution() {
|
||||
let response = error_response(json!(1), "execute_query", Some("s1".into()), anyhow!("connection reset"));
|
||||
let data = response.error.unwrap().data.unwrap();
|
||||
assert_eq!(data.category, "connection");
|
||||
assert_eq!(data.session_disposition, "quarantine");
|
||||
assert!(!data.retryable);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tdengine_uses_database_when_schema_is_empty() {
|
||||
assert_eq!(database_scope(&json!({"database": "power", "schema": ""})), "power");
|
||||
assert_eq!(database_scope(&json!({"database": "power", "schema": "metrics"})), "metrics");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn metadata_constraints_ignore_protocol_session_fields() {
|
||||
let constraints: crate::model::MetadataListConstraints = decode(&json!({
|
||||
"agentSessionId": "session-1",
|
||||
"filter": "meter",
|
||||
"limit": 10,
|
||||
}))
|
||||
.unwrap();
|
||||
assert_eq!(constraints.filter, "meter");
|
||||
assert_eq!(constraints.limit, 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn query_options_accept_protocol_camel_case_fields() {
|
||||
let options: crate::model::QueryOptions =
|
||||
decode(&json!({"sql": "SELECT 1", "pageSize": 25, "maxRows": 100})).unwrap();
|
||||
assert_eq!(options.page_size, 25);
|
||||
assert_eq!(options.max_rows, 100);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn completion_request_accepts_null_optional_fields() {
|
||||
let request: crate::model::CompletionAssistantRequest = decode(&json!({
|
||||
"connection_id": "connection-1",
|
||||
"database": "power",
|
||||
"schema": null,
|
||||
"object_kinds": ["table"],
|
||||
"max_results": null,
|
||||
"parent_schema": null,
|
||||
"parent_name": null,
|
||||
"match_mode": null,
|
||||
}))
|
||||
.unwrap();
|
||||
assert_eq!(request.schema, "");
|
||||
assert_eq!(request.max_results, 0);
|
||||
assert_eq!(request.parent_schema, "");
|
||||
assert_eq!(request.parent_name, "");
|
||||
assert_eq!(request.match_mode, "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancel_session_cancels_the_active_operation() {
|
||||
let slot = Arc::new(Semaphore::new(1)).try_acquire_owned().unwrap();
|
||||
let session = AgentSession {
|
||||
driver: Mutex::new(TdengineDriver::new()),
|
||||
active: StdMutex::new(None),
|
||||
next_operation_id: AtomicU64::new(1),
|
||||
_slot: slot,
|
||||
};
|
||||
let (token, operation) = session.begin_operation();
|
||||
session.cancel();
|
||||
assert!(token.is_cancelled());
|
||||
drop(operation);
|
||||
assert!(session.active.lock().unwrap().is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn session_slots_bound_connects_in_flight_and_open_sessions() {
|
||||
let runtime = RuntimeServer::new();
|
||||
let permits = (0..MAX_AGENT_SESSIONS)
|
||||
.map(|_| runtime.session_slots.clone().try_acquire_owned().unwrap())
|
||||
.collect::<Vec<_>>();
|
||||
assert!(runtime.session_slots.clone().try_acquire_owned().is_err());
|
||||
drop(permits);
|
||||
assert!(runtime.session_slots.clone().try_acquire_owned().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_slots_bound_non_control_work_without_blocking_cancellation() {
|
||||
let runtime = RuntimeServer::new();
|
||||
let permits = (0..MAX_CONCURRENT_REQUESTS)
|
||||
.map(|_| runtime.request_slots.clone().try_acquire_owned().unwrap())
|
||||
.collect::<Vec<_>>();
|
||||
assert!(runtime.request_slots.clone().try_acquire_owned().is_err());
|
||||
assert!(is_capacity_exempt("cancel_session"));
|
||||
assert!(is_capacity_exempt("close_session"));
|
||||
assert!(!is_capacity_exempt("execute_query"));
|
||||
drop(permits);
|
||||
}
|
||||
}
|
||||
-178
@@ -1,178 +0,0 @@
|
||||
package com.dbx.agent.tdengine;
|
||||
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.test.TestSupport;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
import java.lang.reflect.InvocationHandler;
|
||||
import java.lang.reflect.Method;
|
||||
import java.lang.reflect.Proxy;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.sql.Connection;
|
||||
import java.sql.ResultSet;
|
||||
import java.sql.ResultSetMetaData;
|
||||
import java.sql.Statement;
|
||||
import java.sql.Types;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertFalse;
|
||||
import static org.junit.jupiter.api.Assertions.assertTrue;
|
||||
|
||||
class TDengineAgentPagingTest {
|
||||
@Test
|
||||
void executeQueryPageDecodesBinaryVarcharValuesAsStrings() {
|
||||
TDengineAgent agent = new TDengineAgent();
|
||||
TestSupport.setPrivateConnection(agent, binaryVarcharConnection());
|
||||
JsonRpcServer server = new JsonRpcServer(agent);
|
||||
|
||||
String response = handleRequest(server, """
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "execute_query_page",
|
||||
"params": {
|
||||
"sql": "SELECT v FROM sample",
|
||||
"schema": "power",
|
||||
"pageSize": 100
|
||||
}
|
||||
}
|
||||
""");
|
||||
|
||||
assertFalse(response.contains("\"error\""), response);
|
||||
assertTrue(response.contains("\"hello\""), response);
|
||||
assertFalse(response.contains("[104,101,108,108,111]"), response);
|
||||
}
|
||||
|
||||
private static Connection binaryVarcharConnection() {
|
||||
Statement statement = binaryVarcharStatement();
|
||||
return proxy(Connection.class, (method, args) -> {
|
||||
String name = method.getName();
|
||||
if ("createStatement".equals(name)) {
|
||||
return statement;
|
||||
}
|
||||
if ("getAutoCommit".equals(name)) {
|
||||
return true;
|
||||
}
|
||||
if ("setAutoCommit".equals(name) || "commit".equals(name) || "rollback".equals(name) || "close".equals(name)) {
|
||||
return null;
|
||||
}
|
||||
if ("isClosed".equals(name)) {
|
||||
return false;
|
||||
}
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
}
|
||||
|
||||
private static Statement binaryVarcharStatement() {
|
||||
ResultSet resultSet = binaryVarcharResultSet();
|
||||
return proxy(Statement.class, (method, args) -> {
|
||||
String name = method.getName();
|
||||
if ("execute".equals(name)) {
|
||||
return true;
|
||||
}
|
||||
if ("getResultSet".equals(name)) {
|
||||
return resultSet;
|
||||
}
|
||||
if ("getUpdateCount".equals(name)) {
|
||||
return 0;
|
||||
}
|
||||
if ("setMaxRows".equals(name) || "setFetchSize".equals(name) || "close".equals(name)) {
|
||||
return null;
|
||||
}
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
}
|
||||
|
||||
private static ResultSet binaryVarcharResultSet() {
|
||||
final int[] index = {-1};
|
||||
ResultSetMetaData metadata = binaryVarcharMetadata();
|
||||
byte[] value = "hello".getBytes(StandardCharsets.UTF_8);
|
||||
return proxy(ResultSet.class, (method, args) -> {
|
||||
String name = method.getName();
|
||||
if ("next".equals(name)) {
|
||||
index[0] += 1;
|
||||
return index[0] == 0;
|
||||
}
|
||||
if ("getMetaData".equals(name)) {
|
||||
return metadata;
|
||||
}
|
||||
if ("getObject".equals(name)) {
|
||||
return value;
|
||||
}
|
||||
if ("wasNull".equals(name)) {
|
||||
return false;
|
||||
}
|
||||
if ("close".equals(name)) {
|
||||
return null;
|
||||
}
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
}
|
||||
|
||||
private static ResultSetMetaData binaryVarcharMetadata() {
|
||||
return proxy(ResultSetMetaData.class, (method, args) -> {
|
||||
String name = method.getName();
|
||||
if ("getColumnCount".equals(name)) {
|
||||
return 1;
|
||||
}
|
||||
if ("getColumnLabel".equals(name)) {
|
||||
return "v";
|
||||
}
|
||||
if ("getColumnType".equals(name)) {
|
||||
return Types.VARCHAR;
|
||||
}
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
}
|
||||
|
||||
private static String handleRequest(JsonRpcServer server, String request) {
|
||||
try {
|
||||
Method method = JsonRpcServer.class.getDeclaredMethod("handleRequest", String.class);
|
||||
method.setAccessible(true);
|
||||
return (String) method.invoke(server, request);
|
||||
} catch (Exception e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
}
|
||||
|
||||
private static <T> T proxy(Class<T> type, MethodHandler handler) {
|
||||
InvocationHandler invocationHandler = new InvocationHandler() {
|
||||
@Override
|
||||
public Object invoke(Object proxy, Method method, Object[] args) {
|
||||
return handler.handle(method, args);
|
||||
}
|
||||
};
|
||||
return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class<?>[]{type}, invocationHandler));
|
||||
}
|
||||
|
||||
private static Object defaultValue(Class<?> type) {
|
||||
if (Boolean.TYPE.equals(type)) {
|
||||
return false;
|
||||
}
|
||||
if (Byte.TYPE.equals(type)) {
|
||||
return (byte) 0;
|
||||
}
|
||||
if (Short.TYPE.equals(type)) {
|
||||
return (short) 0;
|
||||
}
|
||||
if (Integer.TYPE.equals(type)) {
|
||||
return 0;
|
||||
}
|
||||
if (Long.TYPE.equals(type)) {
|
||||
return 0L;
|
||||
}
|
||||
if (Float.TYPE.equals(type)) {
|
||||
return 0f;
|
||||
}
|
||||
if (Double.TYPE.equals(type)) {
|
||||
return 0.0d;
|
||||
}
|
||||
if (Character.TYPE.equals(type)) {
|
||||
return '\0';
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
private interface MethodHandler {
|
||||
Object handle(Method method, Object[] args);
|
||||
}
|
||||
}
|
||||
@@ -1,584 +0,0 @@
|
||||
package com.dbx.agent.tdengine;
|
||||
|
||||
import com.dbx.agent.ColumnInfo;
|
||||
import com.dbx.agent.ConnectParams;
|
||||
import com.dbx.agent.DatabaseAgent;
|
||||
import com.dbx.agent.ExecuteQueryOptions;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.TableInfo;
|
||||
import com.dbx.agent.test.JdbcAgentFake;
|
||||
import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest;
|
||||
import com.dbx.agent.test.JdbcMetadataSqlFake;
|
||||
import com.dbx.agent.test.TestSupport;
|
||||
import com.taosdata.jdbc.TSDBDriver;
|
||||
import com.taosdata.jdbc.rs.RestfulConnection;
|
||||
import com.zaxxer.hikari.HikariConfig;
|
||||
import com.zaxxer.hikari.HikariDataSource;
|
||||
import java.lang.reflect.InvocationHandler;
|
||||
import java.lang.reflect.Proxy;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.sql.Connection;
|
||||
import java.sql.ResultSet;
|
||||
import java.sql.Timestamp;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Arrays;
|
||||
import java.util.List;
|
||||
import java.util.Properties;
|
||||
import javax.sql.DataSource;
|
||||
import org.junit.jupiter.api.Assertions;
|
||||
import org.junit.jupiter.api.Test;
|
||||
|
||||
class TDengineAgentExecutionTest extends JdbcFakeExecutionBehaviorTest {
|
||||
@Override
|
||||
protected DatabaseAgent createAgent() {
|
||||
return new TDengineAgent();
|
||||
}
|
||||
|
||||
@Override
|
||||
protected String resultSetSql() {
|
||||
return "SHOW DATABASES";
|
||||
}
|
||||
}
|
||||
|
||||
class TDengineAgentMetadataTest {
|
||||
@Test
|
||||
void buildsWebsocketJdbcUrlWithDefaultPortAndDatabase() {
|
||||
String url = TDengineJdbcUrl.from(
|
||||
new ConnectParams("127.0.0.1", 0, "meters", "root", "taosdata", "", "", false)
|
||||
);
|
||||
|
||||
Assertions.assertEquals("jdbc:TAOS-WS://127.0.0.1:6041/meters", url);
|
||||
}
|
||||
|
||||
@Test
|
||||
void preservesCustomWebsocketPortAndUrlParams() {
|
||||
String url = TDengineJdbcUrl.from(
|
||||
new ConnectParams("td.local", 6042, "", "", "", "timezone=UTC&charset=UTF-8", "", false)
|
||||
);
|
||||
|
||||
Assertions.assertEquals("jdbc:TAOS-WS://td.local:6042/?timezone=UTC&charset=UTF-8", url);
|
||||
}
|
||||
|
||||
@Test
|
||||
void supportsRestTransportViaCompatibilityUrlParam() {
|
||||
String url = TDengineJdbcUrl.from(
|
||||
new ConnectParams("127.0.0.1", 0, "testdb", "", "", "dbx.transport=rest&charset=UTF-8", "", false)
|
||||
);
|
||||
|
||||
Assertions.assertEquals("jdbc:TAOS-RS://127.0.0.1:6041/testdb?charset=UTF-8", url);
|
||||
}
|
||||
|
||||
@Test
|
||||
void stripsTransportControlParamFromJdbcQuery() {
|
||||
String url = TDengineJdbcUrl.from(
|
||||
new ConnectParams("127.0.0.1", 6041, "", "", "", "transport=ws&timezone=UTC", "", false)
|
||||
);
|
||||
|
||||
Assertions.assertEquals("jdbc:TAOS-WS://127.0.0.1:6041/?timezone=UTC", url);
|
||||
}
|
||||
|
||||
@Test
|
||||
void usesTdengineShowStatementsForMetadata() {
|
||||
TDengineAgent agent = new TDengineAgent();
|
||||
TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection());
|
||||
|
||||
agent.listDatabases();
|
||||
agent.listTables("power");
|
||||
agent.getColumns("power", "meters");
|
||||
|
||||
Assertions.assertEquals("SHOW DATABASES", JdbcMetadataSqlFake.statements.get(0));
|
||||
Assertions.assertEquals("SHOW `power`.STABLES", JdbcMetadataSqlFake.statements.get(1));
|
||||
Assertions.assertEquals("SHOW `power`.TABLES", JdbcMetadataSqlFake.statements.get(2));
|
||||
Assertions.assertFalse(JdbcMetadataSqlFake.statements.stream().anyMatch(sql -> sql.contains("information_schema")));
|
||||
Assertions.assertTrue(JdbcMetadataSqlFake.statements.contains("DESCRIBE `power`.`meters`"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void showMetadataStopsReadingAtTheRequestedPageBoundary() {
|
||||
List<String> statements = new ArrayList<>();
|
||||
List<Integer> maxRows = new ArrayList<>();
|
||||
List<String> rowReads = new ArrayList<>();
|
||||
TDengineAgent agent = new TDengineAgent();
|
||||
TestSupport.setPrivateConnection(agent, showMetadataConnection(statements, maxRows, rowReads));
|
||||
|
||||
List<TableInfo> firstPage = agent.listTables(
|
||||
"dbx_alpha",
|
||||
new MetadataListConstraints(null, 3, null, List.of("TABLE"))
|
||||
);
|
||||
List<TableInfo> secondPage = agent.listTables(
|
||||
"dbx_alpha",
|
||||
new MetadataListConstraints(null, 2, 3, List.of("TABLE"))
|
||||
);
|
||||
|
||||
Assertions.assertEquals(List.of("meters", "weather", "device_a"), firstPage.stream().map(TableInfo::getName).toList());
|
||||
Assertions.assertEquals(List.of("device_b", "standalone"), secondPage.stream().map(TableInfo::getName).toList());
|
||||
Assertions.assertEquals("meters", firstPage.get(2).getParent_name());
|
||||
Assertions.assertEquals("meters", secondPage.get(0).getParent_name());
|
||||
Assertions.assertTrue(maxRows.isEmpty());
|
||||
Assertions.assertEquals(
|
||||
List.of(
|
||||
"SHOW `dbx_alpha`.STABLES",
|
||||
"SHOW `dbx_alpha`.TABLES",
|
||||
"SHOW `dbx_alpha`.STABLES",
|
||||
"SHOW `dbx_alpha`.TABLES"
|
||||
),
|
||||
statements
|
||||
);
|
||||
Assertions.assertEquals(4, rowReads.stream().filter("STABLE"::equals).count());
|
||||
Assertions.assertEquals(4, rowReads.stream().filter("TABLE"::equals).count());
|
||||
}
|
||||
|
||||
@Test
|
||||
void showMetadataCachesOnlyUnboundedResults() {
|
||||
List<String> statements = new ArrayList<>();
|
||||
List<Integer> maxRows = new ArrayList<>();
|
||||
TDengineAgent agent = new TDengineAgent();
|
||||
TestSupport.setPrivateConnection(agent, showMetadataConnection(statements, maxRows));
|
||||
|
||||
agent.listTables("dbx_alpha");
|
||||
agent.listTables("dbx_alpha");
|
||||
|
||||
Assertions.assertEquals(
|
||||
List.of(
|
||||
"SHOW `dbx_alpha`.STABLES",
|
||||
"SHOW `dbx_alpha`.TABLES"
|
||||
),
|
||||
statements
|
||||
);
|
||||
}
|
||||
|
||||
@Test
|
||||
void showMetadataFiltersLocallyBeforeApplyingRowLimit() {
|
||||
List<String> statements = new ArrayList<>();
|
||||
List<Integer> maxRows = new ArrayList<>();
|
||||
TDengineAgent agent = new TDengineAgent();
|
||||
TestSupport.setPrivateConnection(agent, showMetadataConnection(statements, maxRows));
|
||||
|
||||
List<TableInfo> tables = agent.listTables(
|
||||
"dbx_alpha",
|
||||
new MetadataListConstraints("stand", 1, null, List.of("TABLE"))
|
||||
);
|
||||
|
||||
Assertions.assertEquals(List.of("standalone"), tables.stream().map(TableInfo::getName).toList());
|
||||
Assertions.assertTrue(maxRows.isEmpty());
|
||||
Assertions.assertEquals(
|
||||
List.of(
|
||||
"SHOW `dbx_alpha`.STABLES",
|
||||
"SHOW `dbx_alpha`.TABLES"
|
||||
),
|
||||
statements
|
||||
);
|
||||
}
|
||||
|
||||
@Test
|
||||
void showMetadataFallsBackToLocalFilteringForLongLikePatterns() {
|
||||
List<String> statements = new ArrayList<>();
|
||||
List<Integer> maxRows = new ArrayList<>();
|
||||
TDengineAgent agent = new TDengineAgent();
|
||||
TestSupport.setPrivateConnection(agent, showMetadataConnection(statements, maxRows));
|
||||
|
||||
agent.listTables(
|
||||
"dbx_alpha",
|
||||
new MetadataListConstraints("a".repeat(50), 1, null, List.of("TABLE"))
|
||||
);
|
||||
|
||||
Assertions.assertTrue(maxRows.isEmpty());
|
||||
Assertions.assertEquals(
|
||||
List.of(
|
||||
"SHOW `dbx_alpha`.STABLES",
|
||||
"SHOW `dbx_alpha`.TABLES"
|
||||
),
|
||||
statements
|
||||
);
|
||||
}
|
||||
|
||||
@Test
|
||||
void showMetadataDoesNotLoadTableComments() {
|
||||
List<String> statements = new ArrayList<>();
|
||||
List<Integer> maxRows = new ArrayList<>();
|
||||
TDengineAgent agent = new TDengineAgent();
|
||||
TestSupport.setPrivateConnection(agent, showMetadataConnection(statements, maxRows));
|
||||
|
||||
List<TableInfo> tables = agent.listTables("dbx_alpha");
|
||||
|
||||
Assertions.assertNull(tables.get(0).getComment());
|
||||
Assertions.assertEquals(
|
||||
List.of(
|
||||
"SHOW `dbx_alpha`.STABLES",
|
||||
"SHOW `dbx_alpha`.TABLES"
|
||||
),
|
||||
statements
|
||||
);
|
||||
}
|
||||
|
||||
@Test
|
||||
void showMetadataFiltersByNameWithoutInformationSchema() {
|
||||
List<String> statements = new ArrayList<>();
|
||||
List<Integer> maxRows = new ArrayList<>();
|
||||
TDengineAgent agent = new TDengineAgent();
|
||||
TestSupport.setPrivateConnection(agent, showMetadataConnection(statements, maxRows));
|
||||
|
||||
List<TableInfo> tables = agent.listTables(
|
||||
"dbx_alpha",
|
||||
new MetadataListConstraints("device_b", 1, null, List.of("TABLE"))
|
||||
);
|
||||
|
||||
Assertions.assertEquals(List.of("device_b"), tables.stream().map(TableInfo::getName).toList());
|
||||
Assertions.assertTrue(maxRows.isEmpty());
|
||||
}
|
||||
|
||||
@Test
|
||||
void sortsSupertablesBeforeTheirChildTables() {
|
||||
List<TableInfo> tables = new ArrayList<>(Arrays.asList(
|
||||
new TableInfo("device_b", "TABLE", null, null, "meters"),
|
||||
new TableInfo("standalone", "TABLE", null),
|
||||
new TableInfo("meters", "STABLE", null),
|
||||
new TableInfo("device_a", "TABLE", null, null, "meters")
|
||||
));
|
||||
|
||||
TDengineAgent.sortTablesForHierarchy(tables);
|
||||
|
||||
Assertions.assertEquals(
|
||||
Arrays.asList("meters", "device_a", "device_b", "standalone"),
|
||||
tables.stream().map(TableInfo::getName).toList()
|
||||
);
|
||||
Assertions.assertEquals("meters", tables.get(1).getParent_name());
|
||||
}
|
||||
|
||||
@Test
|
||||
void marksTimestampAndCompositeKeyDescribeColumnsAsPrimaryKeys() throws Exception {
|
||||
ResultSet resultSet = describeResultSet(new String[][] {
|
||||
{"ts", "TIMESTAMP", "8", ""},
|
||||
{"seq", "INT", "4", "COMPOSITE KEY"},
|
||||
{"voltage", "FLOAT", "4", ""},
|
||||
{"site", "VARCHAR(32)", "32", "TAG"}
|
||||
});
|
||||
|
||||
List<ColumnInfo> columns = TDengineAgent.readDescribeColumns(resultSet);
|
||||
|
||||
Assertions.assertTrue(columns.get(0).getIs_primary_key());
|
||||
Assertions.assertFalse(columns.get(0).getIs_nullable());
|
||||
Assertions.assertTrue(columns.get(1).getIs_primary_key());
|
||||
Assertions.assertFalse(columns.get(1).getIs_nullable());
|
||||
Assertions.assertFalse(columns.get(2).getIs_primary_key());
|
||||
Assertions.assertFalse(columns.get(3).getIs_primary_key());
|
||||
Assertions.assertEquals("TAG", columns.get(3).getExtra());
|
||||
}
|
||||
|
||||
@Test
|
||||
void doesNotExposeDatabasesAsSchemas() {
|
||||
TDengineAgent agent = new TDengineAgent();
|
||||
TestSupport.setPrivateConnection(agent, JdbcMetadataSqlFake.connection());
|
||||
|
||||
Assertions.assertTrue(agent.listSchemas().isEmpty());
|
||||
Assertions.assertTrue(JdbcMetadataSqlFake.statements.isEmpty());
|
||||
}
|
||||
|
||||
private static Connection showMetadataConnection(List<String> statements, List<Integer> maxRows) {
|
||||
return showMetadataConnection(statements, maxRows, null);
|
||||
}
|
||||
|
||||
private static Connection showMetadataConnection(
|
||||
List<String> statements,
|
||||
List<Integer> maxRows,
|
||||
List<String> rowReads
|
||||
) {
|
||||
return proxy(Connection.class, (proxy, method, args) -> {
|
||||
String name = method.getName();
|
||||
if ("createStatement".equals(name)) {
|
||||
return showMetadataStatement(statements, maxRows, rowReads);
|
||||
}
|
||||
if ("isClosed".equals(name)) return false;
|
||||
if ("close".equals(name)) return null;
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
}
|
||||
|
||||
private static java.sql.Statement showMetadataStatement(
|
||||
List<String> statements,
|
||||
List<Integer> maxRows,
|
||||
List<String> rowReads
|
||||
) {
|
||||
int[] activeMaxRows = {0};
|
||||
return proxy(java.sql.Statement.class, (proxy, method, args) -> {
|
||||
String name = method.getName();
|
||||
if ("setMaxRows".equals(name)) {
|
||||
activeMaxRows[0] = (Integer) args[0];
|
||||
maxRows.add(activeMaxRows[0]);
|
||||
return null;
|
||||
}
|
||||
if ("executeQuery".equals(name)) {
|
||||
String sql = (String) args[0];
|
||||
statements.add(sql);
|
||||
if (sql.endsWith("STABLES")) {
|
||||
return showTableResultSet(
|
||||
List.of(new String[] {"meters"}, new String[] {"weather"}),
|
||||
activeMaxRows[0],
|
||||
"STABLE",
|
||||
rowReads
|
||||
);
|
||||
}
|
||||
if (sql.endsWith("TABLES")) {
|
||||
return showTableResultSet(
|
||||
List.of(
|
||||
new String[] {"device_a", "", "", "meters"},
|
||||
new String[] {"device_b", "", "", "meters"},
|
||||
new String[] {"standalone", "", "", ""}
|
||||
),
|
||||
activeMaxRows[0],
|
||||
"TABLE",
|
||||
rowReads
|
||||
);
|
||||
}
|
||||
}
|
||||
if ("close".equals(name)) return null;
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
}
|
||||
|
||||
private static ResultSet showTableResultSet(
|
||||
List<String[]> rows,
|
||||
int maxRows,
|
||||
String tableType,
|
||||
List<String> rowReads
|
||||
) {
|
||||
List<String[]> limitedRows = maxRows > 0 ? rows.subList(0, Math.min(rows.size(), maxRows)) : rows;
|
||||
int[] index = {-1};
|
||||
return proxy(ResultSet.class, (proxy, method, args) -> {
|
||||
String name = method.getName();
|
||||
if ("next".equals(name)) {
|
||||
index[0] += 1;
|
||||
boolean hasNext = index[0] < limitedRows.size();
|
||||
if (hasNext && rowReads != null) {
|
||||
rowReads.add(tableType);
|
||||
}
|
||||
return hasNext;
|
||||
}
|
||||
if ("getString".equals(name)) {
|
||||
int column = (Integer) args[0];
|
||||
String[] row = limitedRows.get(index[0]);
|
||||
return column <= row.length ? row[column - 1] : null;
|
||||
}
|
||||
if ("isClosed".equals(name)) return false;
|
||||
if ("close".equals(name)) return null;
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
}
|
||||
|
||||
private static <T> T proxy(Class<T> type, InvocationHandler handler) {
|
||||
return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class<?>[]{type}, handler));
|
||||
}
|
||||
|
||||
private static Object defaultValue(Class<?> type) {
|
||||
if (Boolean.TYPE.equals(type)) return false;
|
||||
if (Integer.TYPE.equals(type)) return 0;
|
||||
if (Long.TYPE.equals(type)) return 0L;
|
||||
if (Double.TYPE.equals(type)) return 0.0d;
|
||||
if (Float.TYPE.equals(type)) return 0.0f;
|
||||
if (Short.TYPE.equals(type)) return (short) 0;
|
||||
if (Byte.TYPE.equals(type)) return (byte) 0;
|
||||
if (Character.TYPE.equals(type)) return '\0';
|
||||
return null;
|
||||
}
|
||||
|
||||
@Test
|
||||
void setsDatabaseBeforeExecutionWhenSchemaIsProvided() {
|
||||
TDengineAgent agent = new TDengineAgent();
|
||||
TestSupport.setPrivateConnection(agent, JdbcAgentFake.connection());
|
||||
|
||||
agent.executeQuery("SELECT 1", "power", new ExecuteQueryOptions());
|
||||
|
||||
// JdbcSchemaSwitcher executes the USE statement (recorded as "execute")
|
||||
// before the query runs, so the database is switched prior to execution.
|
||||
Assertions.assertEquals(
|
||||
Arrays.asList("execute", "setMaxRows:10001", "execute"),
|
||||
JdbcAgentFake.calls
|
||||
);
|
||||
}
|
||||
|
||||
@Test
|
||||
void restExecutionSwitchesCatalogWithoutGeneratingUseSql() throws Exception {
|
||||
try (RestfulConnection connection = restfulConnection()) {
|
||||
String executionSchema = TDengineAgent.prepareExecutionSchema(connection, " power-data ");
|
||||
|
||||
Assertions.assertNull(executionSchema);
|
||||
Assertions.assertEquals("power-data", connection.getCatalog());
|
||||
Assertions.assertEquals(
|
||||
"power-data",
|
||||
connection.getClientInfo(TSDBDriver.PROPERTY_KEY_DBNAME)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void restExecutionSupportsPooledWrappersAndRestoresOriginalDatabase() throws Exception {
|
||||
Properties properties = new Properties();
|
||||
properties.setProperty(TSDBDriver.PROPERTY_KEY_DBNAME, "original");
|
||||
try (RestfulConnection physical = restfulConnection(properties)) {
|
||||
physical.setCatalog("original");
|
||||
Connection pooled = wrapperConnection(physical);
|
||||
|
||||
Assertions.assertNull(TDengineAgent.prepareExecutionSchema(pooled, "power"));
|
||||
Assertions.assertEquals("power", pooled.getCatalog());
|
||||
Assertions.assertEquals("power", pooled.getClientInfo(TSDBDriver.PROPERTY_KEY_DBNAME));
|
||||
|
||||
TDengineAgent.restoreRestfulConnectionState(pooled, "original", "original");
|
||||
Assertions.assertEquals("original", pooled.getCatalog());
|
||||
Assertions.assertEquals("original", pooled.getClientInfo(TSDBDriver.PROPERTY_KEY_DBNAME));
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void hikariCanManageRestfulConnectionsWithoutLosingNativeAccess() throws Exception {
|
||||
RestfulConnection physical = restfulConnection();
|
||||
DataSource dataSource = proxy(DataSource.class, (proxy, method, args) -> {
|
||||
if ("getConnection".equals(method.getName())) {
|
||||
return physical;
|
||||
}
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
HikariConfig config = new HikariConfig();
|
||||
config.setPoolName("tdengine-rest-test");
|
||||
config.setDataSource(dataSource);
|
||||
config.setMaximumPoolSize(1);
|
||||
config.setMinimumIdle(0);
|
||||
config.setInitializationFailTimeout(-1L);
|
||||
|
||||
try (HikariDataSource pool = new HikariDataSource(config);
|
||||
Connection connection = pool.getConnection()) {
|
||||
Assertions.assertTrue(connection.isWrapperFor(RestfulConnection.class));
|
||||
Assertions.assertSame(physical, connection.unwrap(RestfulConnection.class));
|
||||
Assertions.assertNull(TDengineAgent.prepareExecutionSchema(connection, "power"));
|
||||
Assertions.assertEquals("power", connection.getCatalog());
|
||||
}
|
||||
}
|
||||
|
||||
@Test
|
||||
void websocketExecutionKeepsSchemaForUseSqlSwitching() throws Exception {
|
||||
Connection connection = JdbcAgentFake.connection();
|
||||
|
||||
Assertions.assertEquals(
|
||||
"power",
|
||||
TDengineAgent.prepareExecutionSchema(connection, "power")
|
||||
);
|
||||
}
|
||||
|
||||
@Test
|
||||
void decodesTdengineByteArrayTextValues() {
|
||||
Assertions.assertEquals(
|
||||
"d1001",
|
||||
TDengineAgent.decodeTdengineValue("d1001".getBytes(StandardCharsets.UTF_8))
|
||||
);
|
||||
}
|
||||
|
||||
@Test
|
||||
void formatsTdengineTimestampsAsSqlLiterals() {
|
||||
Assertions.assertEquals(
|
||||
"2026-05-16 09:35:58.123",
|
||||
TDengineAgent.decodeTdengineValue(Timestamp.valueOf("2026-05-16 09:35:58.123"))
|
||||
);
|
||||
Assertions.assertEquals(
|
||||
"2026-05-16 09:35:58.000",
|
||||
TDengineAgent.decodeTdengineValue(Timestamp.valueOf("2026-05-16 09:35:58"))
|
||||
);
|
||||
Assertions.assertEquals(
|
||||
"2026-05-16 09:35:58.123456",
|
||||
TDengineAgent.decodeTdengineValue(Timestamp.valueOf("2026-05-16 09:35:58.123456"))
|
||||
);
|
||||
Assertions.assertEquals(
|
||||
"2026-05-16 09:35:58.123456789",
|
||||
TDengineAgent.decodeTdengineValue(Timestamp.valueOf("2026-05-16 09:35:58.123456789"))
|
||||
);
|
||||
}
|
||||
|
||||
@Test
|
||||
void unknownTransportValueDefaultsToWebsocketAndKeepsOtherParams() {
|
||||
String url = TDengineJdbcUrl.from(
|
||||
new ConnectParams("127.0.0.1", 6041, "", "", "", "transport=foo&timezone=UTC", "", false)
|
||||
);
|
||||
|
||||
Assertions.assertEquals("jdbc:TAOS-WS://127.0.0.1:6041/?timezone=UTC", url);
|
||||
}
|
||||
|
||||
@Test
|
||||
void sanitizeConnectionStringStripsTransportControlParams() {
|
||||
String sanitized = TDengineJdbcUrl.sanitizeConnectionString(
|
||||
"jdbc:TAOS-WS://127.0.0.1:6041/db?dbx.transport=rest&charset=UTF-8&transport=ws"
|
||||
);
|
||||
|
||||
Assertions.assertEquals("jdbc:TAOS-WS://127.0.0.1:6041/db?charset=UTF-8", sanitized);
|
||||
}
|
||||
|
||||
@Test
|
||||
void sanitizeConnectionStringKeepsFragmentAndNonControlParams() {
|
||||
String sanitized = TDengineJdbcUrl.sanitizeConnectionString(
|
||||
"jdbc:TAOS-RS://127.0.0.1:6041/db?timezone=UTC&dbx.transport=rest#anchor"
|
||||
);
|
||||
|
||||
Assertions.assertEquals("jdbc:TAOS-RS://127.0.0.1:6041/db?timezone=UTC#anchor", sanitized);
|
||||
}
|
||||
|
||||
@Test
|
||||
void enablesRestInformationSchemaPagingForTdengine338AndNewer() {
|
||||
Assertions.assertFalse(TDengineAgent.supportsRestInformationSchemaPagingVersion(null));
|
||||
Assertions.assertFalse(TDengineAgent.supportsRestInformationSchemaPagingVersion("3.3.7.9"));
|
||||
Assertions.assertTrue(TDengineAgent.supportsRestInformationSchemaPagingVersion("3.3.8.0"));
|
||||
Assertions.assertTrue(TDengineAgent.supportsRestInformationSchemaPagingVersion("3.3.8.1"));
|
||||
Assertions.assertTrue(TDengineAgent.supportsRestInformationSchemaPagingVersion("4.0.0.0"));
|
||||
}
|
||||
|
||||
private static ResultSet describeResultSet(String[][] rows) {
|
||||
int[] rowIndex = {-1};
|
||||
return (ResultSet) Proxy.newProxyInstance(
|
||||
ResultSet.class.getClassLoader(),
|
||||
new Class<?>[] {ResultSet.class},
|
||||
(proxy, method, args) -> {
|
||||
if ("next".equals(method.getName())) {
|
||||
rowIndex[0] += 1;
|
||||
return rowIndex[0] < rows.length;
|
||||
}
|
||||
if ("getString".equals(method.getName())) {
|
||||
return rows[rowIndex[0]][((Integer) args[0]) - 1];
|
||||
}
|
||||
if ("close".equals(method.getName())) {
|
||||
return null;
|
||||
}
|
||||
throw new UnsupportedOperationException(method.getName());
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
private static RestfulConnection restfulConnection() {
|
||||
return restfulConnection(new Properties());
|
||||
}
|
||||
|
||||
private static RestfulConnection restfulConnection(Properties properties) {
|
||||
return new RestfulConnection(
|
||||
"127.0.0.1",
|
||||
"6041",
|
||||
properties,
|
||||
"",
|
||||
"jdbc:TAOS-RS://127.0.0.1:6041/",
|
||||
null,
|
||||
false,
|
||||
null,
|
||||
null
|
||||
);
|
||||
}
|
||||
|
||||
private static Connection wrapperConnection(Connection physical) {
|
||||
return proxy(Connection.class, (proxy, method, args) -> {
|
||||
if ("isWrapperFor".equals(method.getName())) {
|
||||
return ((Class<?>) args[0]).isInstance(physical);
|
||||
}
|
||||
if ("unwrap".equals(method.getName())) {
|
||||
return ((Class<?>) args[0]).cast(physical);
|
||||
}
|
||||
try {
|
||||
return method.invoke(physical, args);
|
||||
} catch (java.lang.reflect.InvocationTargetException error) {
|
||||
throw error.getCause();
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
use base64::engine::general_purpose::STANDARD;
|
||||
use base64::Engine;
|
||||
use chrono::{NaiveDateTime, Timelike};
|
||||
use chrono_tz::Tz;
|
||||
use serde_json::{json, Number, Value};
|
||||
use taos::taos_query::common::Timestamp;
|
||||
use taos::BorrowedValue;
|
||||
|
||||
pub fn borrowed_value_to_json(value: BorrowedValue<'_>, timezone: Option<Tz>) -> Value {
|
||||
match value {
|
||||
BorrowedValue::Null(_) => Value::Null,
|
||||
BorrowedValue::Bool(value) => Value::Bool(value),
|
||||
BorrowedValue::TinyInt(value) => json!(value),
|
||||
BorrowedValue::SmallInt(value) => json!(value),
|
||||
BorrowedValue::Int(value) => json!(value),
|
||||
BorrowedValue::BigInt(value) => json!(value),
|
||||
BorrowedValue::UTinyInt(value) => json!(value),
|
||||
BorrowedValue::USmallInt(value) => json!(value),
|
||||
BorrowedValue::UInt(value) => json!(value),
|
||||
BorrowedValue::UBigInt(value) => json!(value),
|
||||
BorrowedValue::Float(value) => finite_number(value as f64),
|
||||
BorrowedValue::Double(value) => finite_number(value),
|
||||
BorrowedValue::VarChar(value) => Value::String(value.to_string()),
|
||||
BorrowedValue::NChar(value) => Value::String(value.into_owned()),
|
||||
BorrowedValue::Timestamp(value) => Value::String(format_timestamp(value, timezone)),
|
||||
BorrowedValue::Json(value) => serde_json::from_slice(&value)
|
||||
.unwrap_or_else(|_| Value::String(String::from_utf8_lossy(&value).into_owned())),
|
||||
BorrowedValue::VarBinary(value) | BorrowedValue::Blob(value) | BorrowedValue::MediumBlob(value) => {
|
||||
binary_json(&value)
|
||||
}
|
||||
BorrowedValue::Geometry(value) => geometry_json(&value),
|
||||
BorrowedValue::Decimal(value) => Value::String(value.to_string()),
|
||||
BorrowedValue::Decimal64(value) => Value::String(value.to_string()),
|
||||
}
|
||||
}
|
||||
|
||||
fn finite_number(value: f64) -> Value {
|
||||
Number::from_f64(value).map(Value::Number).unwrap_or_else(|| Value::String(value.to_string()))
|
||||
}
|
||||
|
||||
fn binary_json(value: &[u8]) -> Value {
|
||||
json!({ "$binary": STANDARD.encode(value) })
|
||||
}
|
||||
|
||||
fn geometry_json(value: &[u8]) -> Value {
|
||||
if let Ok(text) = std::str::from_utf8(value) {
|
||||
let trimmed = text.trim();
|
||||
if looks_like_wkt(trimmed) {
|
||||
return Value::String(trimmed.to_string());
|
||||
}
|
||||
}
|
||||
binary_json(value)
|
||||
}
|
||||
|
||||
fn looks_like_wkt(value: &str) -> bool {
|
||||
let upper = value.to_ascii_uppercase();
|
||||
["POINT", "LINESTRING", "POLYGON", "MULTIPOINT", "MULTILINESTRING", "MULTIPOLYGON", "GEOMETRYCOLLECTION"]
|
||||
.iter()
|
||||
.any(|prefix| upper.starts_with(prefix))
|
||||
}
|
||||
|
||||
fn format_timestamp(timestamp: Timestamp, timezone: Option<Tz>) -> String {
|
||||
let datetime = match timezone {
|
||||
Some(timezone) => timestamp.to_datetime_with_custom_tz(&timezone).naive_local(),
|
||||
None => timestamp.to_naive_datetime(),
|
||||
};
|
||||
format_naive_datetime(datetime, timestamp)
|
||||
}
|
||||
|
||||
fn format_naive_datetime(datetime: NaiveDateTime, timestamp: Timestamp) -> String {
|
||||
let digits = match timestamp {
|
||||
Timestamp::Milliseconds(_) => 3,
|
||||
Timestamp::Microseconds(_) => 6,
|
||||
Timestamp::Nanoseconds(_) => 9,
|
||||
};
|
||||
let mut fraction = format!("{:09}", datetime.nanosecond());
|
||||
fraction.truncate(digits);
|
||||
while fraction.len() > 3 && fraction.ends_with('0') {
|
||||
fraction.pop();
|
||||
}
|
||||
format!("{}.{}", datetime.format("%Y-%m-%d %H:%M:%S"), fraction)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use taos::Ty;
|
||||
|
||||
#[test]
|
||||
fn preserves_timestamp_precision_with_at_least_milliseconds() {
|
||||
assert_eq!(
|
||||
borrowed_value_to_json(
|
||||
BorrowedValue::Timestamp(Timestamp::Milliseconds(1_704_067_200_123)),
|
||||
Some(chrono_tz::UTC)
|
||||
),
|
||||
Value::String("2024-01-01 00:00:00.123".into())
|
||||
);
|
||||
assert_eq!(
|
||||
borrowed_value_to_json(
|
||||
BorrowedValue::Timestamp(Timestamp::Microseconds(1_704_067_200_123_400)),
|
||||
Some(chrono_tz::UTC)
|
||||
),
|
||||
Value::String("2024-01-01 00:00:00.1234".into())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_text_and_binary_values_distinct() {
|
||||
assert_eq!(borrowed_value_to_json(BorrowedValue::VarChar("hello"), None), json!("hello"));
|
||||
assert_eq!(
|
||||
borrowed_value_to_json(BorrowedValue::VarBinary(vec![0, 255].into()), None),
|
||||
json!({"$binary": "AP8="})
|
||||
);
|
||||
assert_eq!(borrowed_value_to_json(BorrowedValue::Null(Ty::Int), None), Value::Null);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_wkt_geometry_as_text() {
|
||||
assert_eq!(
|
||||
borrowed_value_to_json(BorrowedValue::Geometry(b"POINT (1 2)".as_slice().into()), None),
|
||||
json!("POINT (1 2)")
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user