diff --git a/agents/drivers/hive-go/connector.go b/agents/drivers/hive-go/connector.go index c232d08df8..ce745c1827 100644 --- a/agents/drivers/hive-go/connector.go +++ b/agents/drivers/hive-go/connector.go @@ -89,7 +89,7 @@ func newDiscoveryConnector(config connectionConfig) *discoveryConnector { BrowserResponsePort: config.BrowserResponsePort, BrowserResponseTimeout: config.BrowserResponseTimeout, BrowserDisableSSLCheck: config.BrowserDisableSSLCheck, - WaitForNonQueryCompletion: strings.EqualFold(config.DatabaseType, "impala"), + WaitForNonQueryCompletion: waitsForNonQueryCompletion(config.DatabaseType), }) }, driver: &gohive.Driver{}, @@ -98,6 +98,10 @@ func newDiscoveryConnector(config connectionConfig) *discoveryConnector { } } +func waitsForNonQueryCompletion(databaseType string) bool { + return strings.EqualFold(databaseType, "impala") || strings.EqualFold(databaseType, "kyuubi") +} + func gssapiOptionsFromKerberos(config kerberosConfig) gosasl.GSSAPIOptions { return gosasl.GSSAPIOptions{ ConfigPath: config.ConfigPath, diff --git a/agents/drivers/hive-go/connector_test.go b/agents/drivers/hive-go/connector_test.go index 89a2d35310..af9678152a 100644 --- a/agents/drivers/hive-go/connector_test.go +++ b/agents/drivers/hive-go/connector_test.go @@ -114,6 +114,15 @@ func TestNormalizeHiveAuth(t *testing.T) { } } +func TestWaitsForKyuubiAndImpalaNonQueryCompletion(t *testing.T) { + if !waitsForNonQueryCompletion("kyuubi") || !waitsForNonQueryCompletion("KYUUBI") || !waitsForNonQueryCompletion("impala") { + t.Fatal("Kyuubi and Impala must wait for asynchronous non-query operations") + } + if waitsForNonQueryCompletion("hive") { + t.Fatal("Hive must retain its existing non-query behavior") + } +} + func TestKerberosServiceForEndpoint(t *testing.T) { target := endpoint{Host: "hs2.example.com", Port: 10000, Principal: "hive/hs2.example.com@EXAMPLE.COM"} discovered := connectionConfig{Kerberos: kerberosConfig{Service: "hive"}} diff --git a/agents/drivers/hive-go/main.go b/agents/drivers/hive-go/main.go index 7ad269e491..6af04f087c 100644 --- a/agents/drivers/hive-go/main.go +++ b/agents/drivers/hive-go/main.go @@ -431,7 +431,11 @@ func (server *server) dispatch(method string, params map[string]json.RawMessage) case "list_constraints", "list_partitions", "list_subpartitions": return []any{}, false, nil case "get_object_source": - result, err := server.getTableDDL(stringParam(params, "schema"), firstNonEmpty(stringParam(params, "name"), stringParam(params, "table"))) + result, err := server.getObjectSource( + stringParam(params, "schema"), + firstNonEmpty(stringParam(params, "name"), stringParam(params, "table")), + stringParam(params, "object_type"), + ) return result, false, err case "get_table_ddl": result, err := server.getTableDDL(stringParam(params, "schema"), stringParam(params, "table")) diff --git a/agents/drivers/hive-go/metadata.go b/agents/drivers/hive-go/metadata.go index 8fb1ff97a9..bc0dcf1c1b 100644 --- a/agents/drivers/hive-go/metadata.go +++ b/agents/drivers/hive-go/metadata.go @@ -40,6 +40,13 @@ type objectInfo struct { Valid *bool `json:"valid,omitempty"` } +type objectSource struct { + Name string `json:"name"` + ObjectType string `json:"object_type"` + Schema *string `json:"schema"` + Source string `json:"source"` +} + type columnInfo struct { Name string `json:"name"` DataType string `json:"data_type"` @@ -200,7 +207,11 @@ func (server *server) connectionInfo() (map[string]any, error) { productName := "Apache Hive" compatibilityMode := "hive" driverName := "DBX Hive Go Agent" - if strings.EqualFold(server.params.DatabaseType, "impala") || strings.Contains(strings.ToLower(version), "impalad version") { + if strings.EqualFold(server.params.DatabaseType, "kyuubi") { + productName = "Apache Kyuubi" + compatibilityMode = "kyuubi" + driverName = "DBX Kyuubi Go Agent" + } else if strings.EqualFold(server.params.DatabaseType, "impala") || strings.Contains(strings.ToLower(version), "impalad version") { productName = "Apache Impala" compatibilityMode = "impala" driverName = "DBX Impala Go Agent" @@ -573,6 +584,19 @@ func (server *server) getTableDDL(schema, table string) (string, error) { return strings.Join(lines, "\n") + "\n", nil } +func (server *server) getObjectSource(schema, name, objectType string) (objectSource, error) { + source, err := server.getTableDDL(schema, name) + if err != nil { + return objectSource{}, err + } + return objectSource{ + Name: name, + ObjectType: strings.ToUpper(objectType), + Schema: optionalString(schema), + Source: source, + }, nil +} + func (server *server) getExplainInfo(sqlText string) (string, error) { sqlText = trimStatementSQL(sqlText) if sqlText == "" { diff --git a/agents/drivers/hive-go/metadata_test.go b/agents/drivers/hive-go/metadata_test.go index c9bba860b9..a6c114de1c 100644 --- a/agents/drivers/hive-go/metadata_test.go +++ b/agents/drivers/hive-go/metadata_test.go @@ -3,6 +3,7 @@ package main import ( "context" "database/sql/driver" + "encoding/json" "errors" "reflect" "strings" @@ -20,6 +21,77 @@ func TestShowTablesRowName(t *testing.T) { } } +func TestKyuubiConnectionInfoReportsNativeIdentity(t *testing.T) { + behavior := &scriptedBehavior{ + query: func(ctx context.Context, query string) (driver.Rows, error) { + switch query { + case "SELECT VERSION()": + return newScriptedRows(ctx, []string{"version"}, []string{"STRING"}, [][]driver.Value{{"3.5.8"}}), nil + case "SELECT CURRENT_USER()": + return newScriptedRows(ctx, []string{"current_user"}, []string{"STRING"}, [][]driver.Value{{"dbx"}}), nil + default: + return nil, errors.New("unexpected query: " + query) + } + }, + } + server := newScriptedServer(t, behavior) + server.params.DatabaseType = "kyuubi" + server.config.Username = "fallback" + + info, err := server.connectionInfo() + if err != nil { + t.Fatal(err) + } + if info["compatibilityMode"] != "kyuubi" || info["username"] != "dbx" || info["version"] != "3.5.8" { + t.Fatalf("unexpected Kyuubi connection info: %#v", info) + } + databaseInfo, ok := info["databaseInfo"].(map[string]string) + if !ok || databaseInfo["productName"] != "Apache Kyuubi" || databaseInfo["driverName"] != "DBX Kyuubi Go Agent" { + t.Fatalf("unexpected Kyuubi database identity: %#v", info["databaseInfo"]) + } +} + +func TestGetObjectSourceReturnsProtocolObject(t *testing.T) { + behavior := &scriptedBehavior{ + query: func(ctx context.Context, query string) (driver.Rows, error) { + if query != "SHOW CREATE TABLE `dbx_kyuubi_demo`.`high_value_orders`" { + t.Fatalf("unexpected query: %q", query) + } + return newScriptedRows( + ctx, + []string{"createtab_stmt"}, + []string{"STRING"}, + [][]driver.Value{ + {"CREATE VIEW dbx_kyuubi_demo.high_value_orders"}, + {"AS SELECT id, customer, amount FROM dbx_kyuubi_demo.orders WHERE amount >= 50"}, + }, + ), nil + }, + } + server := newScriptedServer(t, behavior) + defer server.disconnect() + + result, _, err := server.dispatch("get_object_source", map[string]json.RawMessage{ + "schema": json.RawMessage(`"dbx_kyuubi_demo"`), + "name": json.RawMessage(`"high_value_orders"`), + "object_type": json.RawMessage(`"VIEW"`), + }) + if err != nil { + t.Fatal(err) + } + source, ok := result.(objectSource) + if !ok { + t.Fatalf("get_object_source returned %T instead of objectSource", result) + } + if source.Name != "high_value_orders" || source.ObjectType != "VIEW" || source.Schema == nil || *source.Schema != "dbx_kyuubi_demo" { + t.Fatalf("unexpected object source metadata: %#v", source) + } + expected := "CREATE VIEW dbx_kyuubi_demo.high_value_orders\nAS SELECT id, customer, amount FROM dbx_kyuubi_demo.orders WHERE amount >= 50\n" + if source.Source != expected { + t.Fatalf("unexpected object source DDL: %q", source.Source) + } +} + func TestListDatabasesUsesShowDatabasesBeforeHiveServerMetadata(t *testing.T) { behavior := &scriptedBehavior{ query: func(ctx context.Context, sql string) (driver.Rows, error) { diff --git a/apps/desktop/public/icons/database/kyuubi.png b/apps/desktop/public/icons/database/kyuubi.png new file mode 100644 index 0000000000..36c5b798c0 Binary files /dev/null and b/apps/desktop/public/icons/database/kyuubi.png differ diff --git a/apps/desktop/src/components/connection/ConnectionDialog.vue b/apps/desktop/src/components/connection/ConnectionDialog.vue index 6bb617f370..1ff7ba37a3 100644 --- a/apps/desktop/src/components/connection/ConnectionDialog.vue +++ b/apps/desktop/src/components/connection/ConnectionDialog.vue @@ -1162,6 +1162,7 @@ const driverProfiles: Record< trino: { type: "trino", port: 8080, user: "", label: "Trino", icon: "trino" }, prestosql: { type: "prestosql", port: 8080, user: "", label: "PrestoSQL", icon: "presto" }, hive: { type: "hive", port: 10000, user: "", label: "Apache Hive", icon: "hive" }, + kyuubi: { type: "kyuubi", port: 10009, user: "", label: "Apache Kyuubi", icon: "kyuubi", urlParams: "auth=NONE" }, impala: { type: "impala", port: 21050, user: "", label: "Apache Impala", icon: "impala", urlParams: "auth=noSasl" }, spark: { type: "spark", port: 10015, user: "", label: "Apache Spark", icon: "spark" }, db2: { type: "db2", port: 50000, user: "db2inst1", label: "IBM DB2", icon: "db2" }, @@ -2416,7 +2417,7 @@ function applyProfile(val: string, preserveConnectionFields = false) { form.value.connection_string = undefined; form.value.url_params = ""; } - resetHiveKerberosFields(profile.type === "hive" || profile.type === "impala" ? form.value : undefined); + resetHiveKerberosFields(profile.type === "hive" || profile.type === "kyuubi" || profile.type === "impala" ? form.value : undefined); } } @@ -2542,7 +2543,7 @@ watch( resetVictoriaMetricsFields(); } resetElasticsearchProxyFields(config.db_type === "elasticsearch" ? config.external_config : undefined); - resetHiveKerberosFields(config.db_type === "hive" || config.db_type === "impala" ? config : undefined); + resetHiveKerberosFields(config.db_type === "hive" || config.db_type === "kyuubi" || config.db_type === "impala" ? config : undefined); resetDamengJvmOptions(config.db_type === "dameng" ? config : undefined); h2ConnectionMode.value = h2ConnectionModeForConfig(config); customColorInput.value = config.color || ""; @@ -2797,6 +2798,7 @@ const iconTypeMap: Record = { trino: "trino", prestosql: "prestosql", hive: "hive", + kyuubi: "kyuubi", impala: "impala", spark: "spark", db2: "db2", @@ -2879,6 +2881,7 @@ const dbOptions: DbOption[] = [ { value: "trino", label: "Trino" }, { value: "prestosql", label: "PrestoSQL" }, { value: "hive", label: "Hive" }, + { value: "kyuubi", label: "Apache Kyuubi" }, { value: "impala", label: "Apache Impala" }, { value: "spark", label: "Apache Spark" }, { value: "db2", label: "DB2" }, @@ -2925,7 +2928,7 @@ const dbCategoryDefinitions: Array<{ { key: "analytics", titleKey: "connection.databaseCategoryAnalytics", - optionValues: ["cloudberry", "clickhouse", "doris", "starrocks", "databend", "selectdb", "databricks", "saphana", "teradata", "vertica", "exasol", "redshift", "snowflake", "trino", "prestosql", "hive", "impala", "spark", "bigquery", "kylin", "dremio"], + optionValues: ["cloudberry", "clickhouse", "doris", "starrocks", "databend", "selectdb", "databricks", "saphana", "teradata", "vertica", "exasol", "redshift", "snowflake", "trino", "prestosql", "hive", "kyuubi", "impala", "spark", "bigquery", "kylin", "dremio"], }, { key: "domestic", @@ -3767,7 +3770,7 @@ function connectionConfigForSubmit(id: string, generatedName = ""): ConnectionCo config.ssl = !!config.ssl || damengSsl.enabled; config.url_params = applyDamengSslUrlParams(config.url_params, config.ssl, damengSsl.sslFilesPath, damengSsl.sslKeystorePassword, damengSsl.sslProtocol); } - if (config.db_type === "hive" || config.db_type === "impala") { + if (config.db_type === "hive" || config.db_type === "kyuubi" || config.db_type === "impala") { if (hiveAuthMode.value === "kerberos" && !hivePrincipal.value.trim()) { throw new Error(t("connection.hiveKerberosPrincipalRequired")); } @@ -7039,7 +7042,7 @@ function openExternalUrl(url: string) {

{{ t("connection.oracleTnsPathHint") }}

-