github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-19 16:16:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(vector): add ChromaDB support

Browse Source
This commit is contained in:
lexmin0412
2026-06-26 01:41:36 +08:00
committed by GitHub
parent 485188322b
commit c2376494d3
43 changed files with 437 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/desktop/public/icons/database/chromadb.svg
+6
View File
@@ -0,0 +1,6 @@
<svg xmlns="http://www.w3.org/2000/svg" width="1em" height="1em" viewBox="0 -46 256 256">
<path d="M0 -46h256v256H0z" fill="none" />
<ellipse cx="170.667" cy="81.92" fill="#ffde2d" rx="85.333" ry="81.92" />
<ellipse cx="85.333" cy="81.92" fill="#327eff" rx="85.333" ry="81.92" />
<path fill="#ff6446" d="M170.667 81.92c0 45.243-38.206 81.92-85.334 81.92V81.92zm-85.334 0C85.333 36.677 123.538 0 170.667 0v81.92z" />
</svg>
Side by Side

After

Width:  |  Height:  |  Size: 427 B

apps/desktop/src/App.vue
+2 -2
View File
@@ -1052,7 +1052,7 @@ async function handleQuickOpenSelect(item: any) {
await connectionStore.loadMongoDatabases(item.connectionId);
} else if (config?.db_type === "elasticsearch") {
await connectionStore.loadElasticsearchIndices(item.connectionId);
} else if (config?.db_type === "qdrant" || config?.db_type === "milvus" || config?.db_type === "weaviate") {
} else if (config?.db_type === "qdrant" || config?.db_type === "milvus" || config?.db_type === "weaviate" || config?.db_type === "chromadb") {
await connectionStore.loadVectorCollections(item.connectionId);
} else if (config?.db_type === "mq") {
await connectionStore.loadMqTenants(item.connectionId);
@@ -1077,7 +1077,7 @@ async function handleQuickOpenSelect(item: any) {
await connectionStore.loadMongoDatabases(item.connectionId);
} else if (config?.db_type === "elasticsearch") {
await connectionStore.loadElasticsearchIndices(item.connectionId);
} else if (config?.db_type === "qdrant" || config?.db_type === "milvus" || config?.db_type === "weaviate") {
} else if (config?.db_type === "qdrant" || config?.db_type === "milvus" || config?.db_type === "weaviate" || config?.db_type === "chromadb") {
await connectionStore.loadVectorCollections(item.connectionId);
} else if (config?.db_type === "mq") {
await connectionStore.loadMqTenants(item.connectionId);
apps/desktop/src/components/connection/ConnectionDialog.vue
+4 -1
View File
@@ -437,6 +437,7 @@ const driverProfiles: Record<
qdrant: { type: "qdrant", port: 6333, user: "", label: "Qdrant", icon: "qdrant" },
milvus: { type: "milvus", port: 19530, user: "root", label: "Milvus", icon: "milvus" },
weaviate: { type: "weaviate", port: 8080, user: "", label: "Weaviate", icon: "weaviate" },
chromadb: { type: "chromadb", port: 8000, user: "", label: "ChromaDB", icon: "chromadb" },
mariadb: { type: "mysql", port: 3306, user: "root", label: "MariaDB", icon: "mariadb" },
tidb: { type: "mysql", port: 4000, user: "root", label: "TiDB", icon: "tidb" },
oceanbase: { type: "mysql", port: 2881, user: "root", label: "OceanBase", icon: "oceanbase" },
@@ -1062,6 +1063,7 @@ const iconTypeMap: Record<string, string> = {
qdrant: "qdrant",
milvus: "milvus",
weaviate: "weaviate",
chromadb: "chromadb",
mariadb: "mariadb",
tidb: "tidb",
oceanbase: "oceanbase",
@@ -1131,6 +1133,7 @@ const dbOptions: DbOption[] = [
{ value: "qdrant", label: "Qdrant" },
{ value: "milvus", label: "Milvus" },
{ value: "weaviate", label: "Weaviate" },
{ value: "chromadb", label: "ChromaDB" },
{ value: "dm", label: "DM (Dameng)" },
{ value: "opengauss", label: "openGauss" },
{ value: "turso", label: "Turso" },
@@ -1237,7 +1240,7 @@ const sqliteExtensionPaths = computed({
form.value.url_params = setSqliteExtensionPaths(form.value.url_params, value);
},
});
const tlsCapableDatabaseTypes = new Set<DatabaseType>(["mysql", "postgres", "redshift", "gaussdb", "kwdb", "opengauss", "questdb", "redis", "etcd", "clickhouse", "elasticsearch", "qdrant", "milvus", "weaviate", "influxdb"]);
const tlsCapableDatabaseTypes = new Set<DatabaseType>(["mysql", "postgres", "redshift", "gaussdb", "kwdb", "opengauss", "questdb", "redis", "etcd", "clickhouse", "elasticsearch", "qdrant", "milvus", "weaviate", "chromadb", "influxdb"]);
const supportsTlsToggle = computed(() => tlsCapableDatabaseTypes.has(form.value.db_type));
const supportsCaCertificatePath = computed(() => form.value.db_type === "clickhouse");
const supportsGenericUrlParams = computed(() => form.value.db_type !== "manticoresearch");
apps/desktop/src/components/diff/DataCompareDialog.vue
+1 -1
View File
@@ -114,7 +114,7 @@ const showModified = ref(true);
let syncPlanRequestId = 0;
const sqlConnections = computed(() => store.connections.filter((connection) => !["redis", "mongodb", "elasticsearch", "qdrant", "milvus", "weaviate", "etcd", "zookeeper", "mq", "nacos"].includes(connection.db_type)));
const sqlConnections = computed(() => store.connections.filter((connection) => !["redis", "mongodb", "elasticsearch", "qdrant", "milvus", "weaviate", "chromadb", "etcd", "zookeeper", "mq", "nacos"].includes(connection.db_type)));
const selectedSourceTableNames = computed(() => sourceTables.value.filter((table) => selectedSourceTables.value.has(table)));
const isBatchCompare = computed(() => selectedSourceTableNames.value.length > 1);
const filteredSourceTables = computed(() => {
apps/desktop/src/components/export/DatabaseExportDialog.vue
+1 -1
View File
@@ -65,7 +65,7 @@ const exportCancelled = ref(false);
const pendingPrefillTable = ref("");
const pendingPrefillTables = ref<string[]>([]);
const sqlConnections = computed(() => store.connections.filter((c) => !["redis", "mongodb", "elasticsearch", "qdrant", "milvus", "weaviate", "etcd", "zookeeper", "mq", "nacos"].includes(c.db_type)));
const sqlConnections = computed(() => store.connections.filter((c) => !["redis", "mongodb", "elasticsearch", "qdrant", "milvus", "weaviate", "chromadb", "etcd", "zookeeper", "mq", "nacos"].includes(c.db_type)));
const canExport = computed(() => connectionId.value && database.value && schema.value && !loadingTables.value && !tableError.value && (tables.value.length === 0 || selectedTables.value.length > 0) && (includeStructure.value || includeData.value || includeObjects.value) && !isExporting.value);
apps/desktop/src/components/icons/DatabaseIcon.vue
+1
View File
@@ -80,6 +80,7 @@ const assetIcons: Record<string, string> = {
qdrant: "qdrant",
milvus: "milvus.png",
weaviate: "weaviate.png",
chromadb: "chromadb",
mq: "pulsar",
pulsar: "pulsar",
nacos: "nacos.png",
apps/desktop/src/components/layout/ContentArea.vue
+21 -2
View File
@@ -62,7 +62,7 @@ import { chartableColumnIndexes } from "@/lib/chartData";
import type { SqlExecutionOverride } from "@/lib/sqlExecutionTarget";
import type { DataGridSortMode } from "@/lib/dataGridSort";
import { useTabScroll } from "@/composables/useTabScroll";
import type { QueryTab, ConnectionConfig, TableInfoTab } from "@/types/database";
import type { QueryTab, ConnectionConfig, TableInfoTab, TreeNode, VectorCollectionMeta } from "@/types/database";
import type { SqlFormatDialect } from "@/lib/sqlFormatter";
type DataGridHandle = {
@@ -177,6 +177,25 @@ const activeTableMeta = computed(() => props.activeTab.tableMeta);
const activeDataTabTableMeta = computed(() => tableMetaForDataTab(props.activeTab));
const activeEffectiveDatabaseType = computed(() => effectiveDatabaseTypeForConnection(props.activeConnection));
function findNodeInTree(nodes: TreeNode[], id: string): TreeNode | undefined {
for (const node of nodes) {
if (node.id === id) return node;
if (node.children) {
const found = findNodeInTree(node.children, id);
if (found) return found;
}
}
return undefined;
}
const activeTabDimension = computed(() => {
const tab = props.activeTab;
if (!tab.connectionId || tab.mode !== "vector") return undefined;
const nodeId = `${tab.connectionId}:__vector_collection:${tab.sql}`;
const meta = findNodeInTree(connectionStore.treeNodes, nodeId)?.meta;
return meta && "dimension" in meta ? (meta as VectorCollectionMeta).dimension : undefined;
});
const activeSqlFormatDialect = computed<SqlFormatDialect>(() => {
switch (activeEffectiveDatabaseType.value) {
case "mysql":
@@ -1068,7 +1087,7 @@ defineExpose({ focusSearch, refreshData, handleModRTarget, requestQueryEditorExe
<!-- Vector mode: Qdrant and Milvus collections -->
<template v-else-if="activeTab.mode === 'vector'">
<div class="flex-1 min-h-0">
<VectorBrowser :key="activeTab.id" :connection-id="activeTab.connectionId" :database="activeTab.database" :collection="activeTab.sql" :database-type="activeEffectiveDatabaseType" />
<VectorBrowser :key="activeTab.id" :connection-id="activeTab.connectionId" :database="activeTab.database" :collection="activeTab.sql" :collection-label="activeTab.title" :database-type="activeEffectiveDatabaseType" :dimension="activeTabDimension" />
</div>
</template>
apps/desktop/src/components/layout/EditorToolbar.vue
+5 -2
View File
@@ -61,7 +61,7 @@ const activeConnectionValue = computed(() => props.activeConnection?.id || "");
const activeSchemaValue = computed(() => props.activeTab.schema || "");
const supportsExplain = computed(() => {
const dbType = props.activeConnection?.db_type;
return dbType !== "redis" && dbType !== "mongodb" && dbType !== "elasticsearch" && dbType !== "qdrant" && dbType !== "milvus" && dbType !== "weaviate" && dbType !== "etcd" && dbType !== "zookeeper" && dbType !== "mq" && dbType !== "nacos";
return dbType !== "redis" && dbType !== "mongodb" && dbType !== "elasticsearch" && dbType !== "qdrant" && dbType !== "milvus" && dbType !== "weaviate" && dbType !== "chromadb" && dbType !== "etcd" && dbType !== "zookeeper" && dbType !== "mq" && dbType !== "nacos";
});
const isSingleDb = computed(() => isSingleDatabase(props.activeConnection?.db_type));
const hasDefaultDatabaseOption = computed(() => activeDatabaseOptions.value.includes(""));
@@ -251,7 +251,10 @@ function connectionById(connectionId: string): ConnectionConfig | undefined {
</template>
</SearchableSelect>
</div>
<div v-if="activeConnection?.db_type !== 'elasticsearch' && activeConnection?.db_type !== 'qdrant' && activeConnection?.db_type !== 'milvus' && activeConnection?.db_type !== 'weaviate' && activeConnection?.db_type !== 'zookeeper' && !isSingleDb" class="flex items-center gap-1">
<div
v-if="activeConnection?.db_type !== 'elasticsearch' && activeConnection?.db_type !== 'qdrant' && activeConnection?.db_type !== 'milvus' && activeConnection?.db_type !== 'weaviate' && activeConnection?.db_type !== 'chromadb' && activeConnection?.db_type !== 'zookeeper' && !isSingleDb"
class="flex items-center gap-1"
>
<Database class="h-3.5 w-3.5 shrink-0" />
<SearchableSelect
:model-value="activeDatabaseValue"
apps/desktop/src/components/sidebar/ConnectionTree.vue
+2 -7
View File
@@ -71,12 +71,7 @@ const simpleObjectParentTypes = new Set<TreeNodeType>(["database", "schema", "li
const simpleObjectChildTypes = new Set<TreeNodeType>(["table", "view", "materialized_view", "procedure", "function", "sequence", "package", "package-body", "load-more"]);
function isSimpleObjectSearchParent(node: TreeNode): boolean {
return (
settingsStore.editorSettings.sidebarObjectDisplay === "simple" &&
simpleObjectParentTypes.has(node.type) &&
node.isExpanded === true &&
!!node.children?.some((child) => simpleObjectChildTypes.has(child.type))
);
return settingsStore.editorSettings.sidebarObjectDisplay === "simple" && simpleObjectParentTypes.has(node.type) && node.isExpanded === true && !!node.children?.some((child) => simpleObjectChildTypes.has(child.type));
}
function collectExpandedObjectSearchTargets(node: TreeNode, tasks: Promise<void>[], refreshedNodeIds?: Set<string>) {
@@ -438,7 +433,7 @@ async function ensureTreeLoadedForTarget(target: ActiveTabSidebarTarget, opts?:
await store.loadMongoDatabases(connId);
} else if (config.db_type === "elasticsearch") {
await store.loadElasticsearchIndices(connId);
} else if (config.db_type === "qdrant" || config.db_type === "milvus" || config.db_type === "weaviate") {
} else if (config.db_type === "qdrant" || config.db_type === "milvus" || config.db_type === "weaviate" || config.db_type === "chromadb") {
await store.loadVectorCollections(connId);
} else if (config.db_type === "mq") {
await store.loadMqTenants(connId, loadOptions);
apps/desktop/src/components/sidebar/TreeItem.vue
+5 -3
View File
@@ -355,6 +355,7 @@ function connectionTooltipScheme(config: Pick<ConnectionConfig, "db_type" | "ssl
case "qdrant":
case "milvus":
case "weaviate":
case "chromadb":
case "rqlite":
case "turso":
case "mq":
@@ -481,7 +482,7 @@ async function toggle() {
await connectionStore.loadMongoDatabases(node.connectionId);
} else if (config?.db_type === "elasticsearch") {
await connectionStore.loadElasticsearchIndices(node.connectionId);
} else if (config?.db_type === "qdrant" || config?.db_type === "milvus" || config?.db_type === "weaviate") {
} else if (config?.db_type === "qdrant" || config?.db_type === "milvus" || config?.db_type === "weaviate" || config?.db_type === "chromadb") {
await connectionStore.loadVectorCollections(node.connectionId);
} else if (config?.db_type === "mq") {
await connectionStore.loadMqTenants(node.connectionId);
@@ -515,8 +516,9 @@ async function toggle() {
const tab = queryStore.createTab(node.connectionId, node.database || "default", node.label, "mongo");
queryStore.updateSql(tab, node.label);
} else if (node.type === "vector-collection" && node.connectionId) {
const collectionRef = node.id.includes("__vector_collection:") ? node.id.split("__vector_collection:").pop() || node.label : node.label;
const tab = queryStore.createTab(node.connectionId, node.database || "default", node.label, "vector");
queryStore.updateSql(tab, node.label);
queryStore.updateSql(tab, collectionRef);
} else if (node.type === "database" && node.connectionId && hasTreeNodeDatabaseContext(node)) {
const config = connectionStore.getConfig(node.connectionId);
const effectiveDbType = effectiveDatabaseTypeForConnection(config);
@@ -3098,7 +3100,7 @@ const nodeIconClass = computed(() => {
const canConfigureVisibleDatabases = computed(() => {
if (props.node.type !== "connection" || !props.node.connectionId) return false;
const dbType = connectionStore.getConfig(props.node.connectionId)?.db_type;
return dbType !== "elasticsearch" && dbType !== "qdrant" && dbType !== "milvus" && dbType !== "weaviate" && dbType !== "etcd" && dbType !== "mq" && dbType !== "nacos";
return dbType !== "elasticsearch" && dbType !== "qdrant" && dbType !== "milvus" && dbType !== "weaviate" && dbType !== "chromadb" && dbType !== "etcd" && dbType !== "mq" && dbType !== "nacos";
});
const canConfigureVisibleSchemas = computed(() => {
apps/desktop/src/components/sql-file/SqlFileExecutionDialog.vue
+1 -1
View File
@@ -49,7 +49,7 @@ const terminalStatus = ref<SqlFileStatus | "idle">("idle");
const terminalError = ref("");
const refreshedTarget = ref(false);
const sqlConnections = computed(() => store.connections.filter((c) => !["redis", "mongodb", "elasticsearch", "qdrant", "milvus", "weaviate", "etcd", "zookeeper", "mq", "nacos"].includes(c.db_type)));
const sqlConnections = computed(() => store.connections.filter((c) => !["redis", "mongodb", "elasticsearch", "qdrant", "milvus", "weaviate", "chromadb", "etcd", "zookeeper", "mq", "nacos"].includes(c.db_type)));
const selectedConnection = computed(() => sqlConnections.value.find((c) => c.id === connectionId.value));
apps/desktop/src/components/transfer/DataTransferDialog.vue
+1 -1
View File
@@ -155,7 +155,7 @@ async function loadTables() {
loadingTables.value = true;
try {
if (isMongoConnection(sourceConnectionId.value)) {
sourceTables.value = await api.mongoListCollections(sourceConnectionId.value, sourceDatabase.value);
sourceTables.value = (await api.mongoListCollections(sourceConnectionId.value, sourceDatabase.value)).map((c) => c.name);
selectedTables.value = new Set(sourceTables.value);
return;
}
apps/desktop/src/components/vector/VectorBrowser.vue
+31 -3
View File
@@ -18,7 +18,9 @@ const props = defineProps<{
connectionId: string;
database: string;
collection: string;
collectionLabel?: string;
databaseType?: DatabaseType;
dimension?: number;
}>();
const loading = ref(false);
@@ -32,8 +34,19 @@ const operationMode = ref<VectorOperationMode>("browse");
const requestText = ref(defaultRequestText(props.databaseType, props.database, props.collection, operationMode.value));
let loadingTimer: ReturnType<typeof setInterval> | undefined;
const productLabel = computed(() => (props.databaseType === "milvus" ? "Milvus" : props.databaseType === "weaviate" ? "Weaviate" : "Qdrant"));
const collectionLabel = computed(() => props.collection || t("vector.collectionFallback"));
const productLabel = computed(() => {
switch (props.databaseType) {
case "milvus":
return "Milvus";
case "weaviate":
return "Weaviate";
case "chromadb":
return "ChromaDB";
default:
return "Qdrant";
}
});
const collectionLabel = computed(() => props.collectionLabel || props.collection || t("vector.collectionFallback"));
const executeLabel = computed(() => (operationMode.value === "browse" ? t("vector.run") : t("vector.apply")));
const operationIcon = computed(() => (operationMode.value === "delete" ? Trash2 : operationMode.value === "upsert" ? Save : Play));
@@ -114,6 +127,18 @@ function defaultRequestText(databaseType: DatabaseType | undefined, database: st
}
return `GET /v1/objects?class=${encodeURIComponent(collectionName)}&limit=100`;
}
if (databaseType === "chromadb") {
const collectionId = collection || "collection-id";
if (mode === "delete") {
return `POST /api/v2/tenants/default_tenant/databases/default_database/collections/${encodeURIComponent(collectionId)}/delete\n${JSON.stringify({ ids: ["id1"] }, null, 2)}`;
}
if (mode === "upsert") {
const dim = props.dimension || 4;
const emb = Array.from({ length: dim }, (_, i) => +(i * 0.1 + 0.1).toFixed(1));
return `POST /api/v2/tenants/default_tenant/databases/default_database/collections/${encodeURIComponent(collectionId)}/upsert\n${JSON.stringify({ ids: ["id1"], embeddings: [emb], documents: ["sample document"], metadatas: [{}] }, null, 2)}`;
}
return `POST /api/v2/tenants/default_tenant/databases/default_database/collections/${encodeURIComponent(collectionId)}/get\n${JSON.stringify({ limit: 100, include: ["documents", "metadatas"] }, null, 2)}`;
}
const collectionPath = pathSegment(collection);
if (mode === "delete") {
return `POST /collections/${collectionPath}/points/delete?wait=true\n${JSON.stringify({ points: [1] }, null, 2)}`;
@@ -238,7 +263,10 @@ function setOperationMode(mode: VectorOperationMode) {
<div class="flex shrink-0 items-center justify-between gap-3 border-b px-3 py-2">
<div class="min-w-0">
<div class="truncate text-sm font-semibold">{{ collectionLabel }}</div>
<div class="truncate text-xs text-muted-foreground">{{ t("vector.productCollection", { product: productLabel }) }}</div>
<div class="truncate text-xs text-muted-foreground">
{{ t("vector.productCollection", { product: productLabel }) }}
<span v-if="dimension != null" class="ml-1.5 inline-flex items-center rounded border bg-muted/50 px-1.5 py-px text-[11px] font-medium text-foreground/80">{{ dimension }}d</span>
</div>
</div>
<div class="flex shrink-0 items-center gap-1.5">
<div class="mr-1 flex h-7 overflow-hidden rounded-md border bg-muted/30 p-0.5">
apps/desktop/src/composables/useNavigationTargets.ts
+1 -1
View File
@@ -25,7 +25,7 @@ async function openTableTarget(target: NavigationTarget, options: { tableInfoTab
connectionStore.activeConnectionId = target.connectionId;
const config = connectionStore.getConfig(target.connectionId);
const tabTitle = target.schema ? `${target.schema}.${target.tableName}` : target.tableName;
if (config?.db_type === "qdrant" || config?.db_type === "milvus" || config?.db_type === "weaviate") {
if (config?.db_type === "qdrant" || config?.db_type === "milvus" || config?.db_type === "weaviate" || config?.db_type === "chromadb") {
await connectionStore.ensureConnected(target.connectionId);
const tabId = queryStore.createTab(target.connectionId, target.database || "default", tabTitle, "vector");
queryStore.updateSql(tabId, target.tableName);
apps/desktop/src/lib/__tests__/sqlStatementRanges.spec.ts
+1
View File
@@ -331,6 +331,7 @@ describe("supportsExecutionTargetPicker", () => {
expect(supportsExecutionTargetPicker("qdrant")).toBe(false);
expect(supportsExecutionTargetPicker("milvus")).toBe(false);
expect(supportsExecutionTargetPicker("weaviate")).toBe(false);
expect(supportsExecutionTargetPicker("chromadb")).toBe(false);
expect(supportsExecutionTargetPicker("etcd")).toBe(false);
expect(supportsExecutionTargetPicker("zookeeper")).toBe(false);
expect(supportsExecutionTargetPicker("mq")).toBe(false);
apps/desktop/src/lib/connectionPresentation.ts
+1
View File
@@ -133,6 +133,7 @@ export function connectionUrlPlaceholder(dbType: DatabaseType): string {
case "qdrant":
case "milvus":
case "weaviate":
case "chromadb":
return "http://user:password@host:port";
case "dameng":
apps/desktop/src/lib/connectionUrl.ts
+2
View File
@@ -43,6 +43,7 @@ const SCHEME_PROFILES: Record<string, ConnectionProfile> = {
qdrant: { type: "qdrant", profile: "qdrant", label: "Qdrant", defaultPort: 6333 },
milvus: { type: "milvus", profile: "milvus", label: "Milvus", defaultPort: 19530 },
weaviate: { type: "weaviate", profile: "weaviate", label: "Weaviate", defaultPort: 8080 },
chromadb: { type: "chromadb", profile: "chromadb", label: "ChromaDB", defaultPort: 8000 },
dm: { type: "dameng", profile: "dm", label: "DM (Dameng)", defaultPort: 5236 },
dameng: { type: "dameng", profile: "dm", label: "DM (Dameng)", defaultPort: 5236 },
gaussdb: { type: "gaussdb", profile: "gaussdb", label: "GaussDB", defaultPort: 5432 },
@@ -66,6 +67,7 @@ const HTTP_SELECTED_PROFILES: Record<string, ConnectionProfile> = {
qdrant: SCHEME_PROFILES.qdrant,
milvus: SCHEME_PROFILES.milvus,
weaviate: SCHEME_PROFILES.weaviate,
chromadb: SCHEME_PROFILES.chromadb,
};
function decodeUrlPart(value: string): string {
apps/desktop/src/lib/connectionVisibleDatabases.ts
+1 -1
View File
@@ -3,7 +3,7 @@ import { filterDatabaseNamesForConnection, normalizeVisibleDatabaseSelection } f
const DRAFT_VISIBLE_DATABASES_PREFIX = "__visible_draft_";
const UNSUPPORTED_VISIBLE_DATABASE_TYPES = new Set<DatabaseType>(["elasticsearch", "qdrant", "milvus", "weaviate", "etcd", "zookeeper"]);
const UNSUPPORTED_VISIBLE_DATABASE_TYPES = new Set<DatabaseType>(["elasticsearch", "qdrant", "milvus", "weaviate", "chromadb", "etcd", "zookeeper"]);
type VisibleDatabaseConnectionFields = Pick<
ConnectionConfig,
apps/desktop/src/lib/http.ts
+5 -3
View File
@@ -30,6 +30,7 @@ import type {
SavedSqlFolder,
SavedSqlLibrary,
} from "@/types/database";
import type { CollectionInfo } from "@/types/database";
import type { SchemaDiffPreparation, SchemaDiffPreparationOptions, TableDiff, FunctionDiff, SequenceDiff, RuleDiff, OwnerDiff } from "@/lib/schemaDiff";
import type { SidebarObjectKind } from "@/lib/databaseObjectCapabilities";
import type { AiConfig, AiTestConnectionResult } from "@/stores/settingsStore";
@@ -1684,7 +1685,7 @@ export async function mongoListDatabases(connectionId: string): Promise<string[]
return post("/api/mongo/list-databases", { connectionId });
}
export async function mongoListCollections(connectionId: string, database: string): Promise<string[]> {
export async function mongoListCollections(connectionId: string, database: string): Promise<CollectionInfo[]> {
return post("/api/mongo/list-collections", { connectionId, database });
}
@@ -1701,10 +1702,11 @@ export async function mongoDropCollection(connectionId: string, database: string
}
export async function elasticsearchListIndices(connectionId: string): Promise<string[]> {
return mongoListCollections(connectionId, "default");
const collections = await mongoListCollections(connectionId, "default");
return collections.map((c) => c.name);
}
export async function vectorListCollections(connectionId: string): Promise<string[]> {
export async function vectorListCollections(connectionId: string): Promise<CollectionInfo[]> {
return mongoListCollections(connectionId, "default");
}
apps/desktop/src/lib/sidebarNodeOrdering.ts
+1 -1
View File
@@ -30,7 +30,7 @@ export function sortSidebarTreeChildrenForParent(parent: Pick<TreeNode, "type">,
const regularChildren = normalized.filter((child) => child.type !== "user-admin" && child.type !== "saved-sql-root");
const withConnectionUtilityOrder = (children: TreeNode[]) => [...savedSqlNodes, ...children, ...userAdminNodes];
if (databaseType === "mongodb" || databaseType === "elasticsearch" || databaseType === "qdrant" || databaseType === "milvus" || databaseType === "weaviate") {
if (databaseType === "mongodb" || databaseType === "elasticsearch" || databaseType === "qdrant" || databaseType === "milvus" || databaseType === "weaviate" || databaseType === "chromadb") {
return withConnectionUtilityOrder(sortByLabel(regularChildren));
}
apps/desktop/src/lib/sqlSemanticDiagnostics.ts
+1 -1
View File
@@ -80,7 +80,7 @@ export function areSqlSemanticDiagnosticsEqual(left: readonly SqlSemanticDiagnos
}
export function shouldRunSqlSemanticDiagnostics(sql: string, cursor: number, options: { databaseType?: DatabaseType } = {}): boolean {
if (options.databaseType === "mongodb" || options.databaseType === "elasticsearch" || options.databaseType === "qdrant" || options.databaseType === "milvus" || options.databaseType === "weaviate" || options.databaseType === "redis") return false;
if (options.databaseType === "mongodb" || options.databaseType === "elasticsearch" || options.databaseType === "qdrant" || options.databaseType === "milvus" || options.databaseType === "weaviate" || options.databaseType === "chromadb" || options.databaseType === "redis") return false;
const context = getSqlCompletionContext(sql, cursor);
if (context.suggestTables || context.exclusiveTableSuggestions || context.exclusiveColumnSuggestions) return false;
if (context.qualifier) return false;
apps/desktop/src/lib/sqlStatementRanges.ts
+2 -1
View File
@@ -11,7 +11,7 @@ export interface SqlTextRange {
sql: string;
}
const NON_SQL_EXECUTION_TARGET_TYPES: ReadonlySet<DatabaseType> = new Set(["mongodb", "elasticsearch", "qdrant", "milvus", "weaviate", "etcd", "zookeeper", "mq", "neo4j"]);
const NON_SQL_EXECUTION_TARGET_TYPES: ReadonlySet<DatabaseType> = new Set(["mongodb", "elasticsearch", "qdrant", "milvus", "weaviate", "chromadb", "etcd", "zookeeper", "mq", "neo4j"]);
export function supportsExecutionTargetPicker(databaseType?: DatabaseType): boolean {
return !!databaseType && (databaseType === "redis" || !NON_SQL_EXECUTION_TARGET_TYPES.has(databaseType));
@@ -89,6 +89,7 @@ const DATABASE_SOFT_STATEMENT_KEYWORDS: Partial<Record<DatabaseType, readonly st
qdrant: [],
milvus: [],
weaviate: [],
chromadb: [],
mq: [],
etcd: [],
zookeeper: [],
apps/desktop/src/lib/tableMetadataCapabilities.ts
+6
View File
@@ -49,6 +49,12 @@ const capabilityByType: Partial<Record<DatabaseType, Partial<TableMetadataCapabi
triggers: false,
ddl: false,
},
chromadb: {
indexes: false,
foreignKeys: false,
triggers: false,
ddl: false,
},
influxdb: {
indexes: false,
foreignKeys: false,
apps/desktop/src/lib/tauri.ts
+5 -3
View File
@@ -31,6 +31,7 @@ import type {
SavedSqlFolder,
SavedSqlLibrary,
} from "@/types/database";
import type { CollectionInfo } from "@/types/database";
import type { SidebarObjectKind } from "@/lib/databaseObjectCapabilities";
import type { AiConfig, AiTestConnectionResult } from "@/stores/settingsStore";
import type { QueryEditability } from "@/lib/sqlAnalysis";
@@ -1394,7 +1395,7 @@ export async function mongoListDatabases(connectionId: string): Promise<string[]
return invoke("mongo_list_databases", { connectionId });
}
export async function mongoListCollections(connectionId: string, database: string): Promise<string[]> {
export async function mongoListCollections(connectionId: string, database: string): Promise<CollectionInfo[]> {
return invoke("mongo_list_collections", { connectionId, database });
}
@@ -1411,10 +1412,11 @@ export async function mongoDropCollection(connectionId: string, database: string
}
export async function elasticsearchListIndices(connectionId: string): Promise<string[]> {
return mongoListCollections(connectionId, "default");
const collections = await mongoListCollections(connectionId, "default");
return collections.map((c) => c.name);
}
export async function vectorListCollections(connectionId: string): Promise<string[]> {
export async function vectorListCollections(connectionId: string): Promise<CollectionInfo[]> {
return mongoListCollections(connectionId, "default");
}
apps/desktop/src/stores/connectionStore.ts
+11 -7
View File
@@ -410,6 +410,7 @@ export const useConnectionStore = defineStore("connection", () => {
qdrant: "Qdrant",
milvus: "Milvus",
weaviate: "Weaviate",
chromadb: "ChromaDB",
doris: "Doris",
starrocks: "StarRocks",
manticoresearch: "Manticore Search",
@@ -1061,7 +1062,7 @@ export const useConnectionStore = defineStore("connection", () => {
await loadMongoDatabases(connectionId);
} else if (config.db_type === "elasticsearch") {
await loadElasticsearchIndices(connectionId);
} else if (config.db_type === "qdrant" || config.db_type === "milvus" || config.db_type === "weaviate") {
} else if (config.db_type === "qdrant" || config.db_type === "milvus" || config.db_type === "weaviate" || config.db_type === "chromadb") {
await loadVectorCollections(connectionId);
} else if (config.db_type === "mq") {
await loadMqTenants(connectionId, { force: true });
@@ -1596,17 +1597,19 @@ export const useConnectionStore = defineStore("connection", () => {
try {
await ensureConnected(connectionId);
const collections = await withMetadataLoadTimeout(connectionId, api.vectorListCollections(connectionId), "vector collections");
const sorted = [...collections].sort((a, b) => a.name.localeCompare(b.name));
setChildren(
node,
withSavedSqlRoot(
connectionId,
sortSidebarNames(collections).map((collection) => ({
id: `${connectionId}:__vector_collection:${collection}`,
label: collection,
sorted.map((info) => ({
id: `${connectionId}:__vector_collection:${info.id}`,
label: info.name,
type: "vector-collection" as const,
connectionId,
database: "default",
isExpanded: false,
meta: info.dimension != null ? { dimension: info.dimension } : undefined,
})),
node,
),
@@ -1628,9 +1631,10 @@ export const useConnectionStore = defineStore("connection", () => {
node.isLoading = true;
try {
const collections = await api.mongoListCollections(connectionId, database);
const names = collections.map((c) => c.name);
setChildren(
node,
sortSidebarNames(collections).map((col) => ({
sortSidebarNames(names).map((col) => ({
id: `${nodeId}:${col}`,
label: col,
type: "mongo-collection" as const,
@@ -2291,7 +2295,7 @@ export const useConnectionStore = defineStore("connection", () => {
await loadMongoDatabases(node.connectionId);
} else if (config?.db_type === "elasticsearch") {
await loadElasticsearchIndices(node.connectionId);
} else if (config?.db_type === "qdrant" || config?.db_type === "milvus" || config?.db_type === "weaviate") {
} else if (config?.db_type === "qdrant" || config?.db_type === "milvus" || config?.db_type === "weaviate" || config?.db_type === "chromadb") {
await loadVectorCollections(node.connectionId);
} else if (config?.db_type === "mq") {
await loadMqTenants(node.connectionId, options);
@@ -2779,7 +2783,7 @@ export const useConnectionStore = defineStore("connection", () => {
if (cached) return cached;
return withCompletionInFlight(`${cacheKey}:mongo-collections`, async () => {
await ensureConnected(connectionId);
const collections = sortSidebarNames(await api.mongoListCollections(connectionId, database));
const collections = sortSidebarNames((await api.mongoListCollections(connectionId, database)).map((c) => c.name));
mongoCompletionCollectionsCache.value[cacheKey] = collections;
evictOldestCacheEntries(mongoCompletionCollectionsCache.value, COMPLETION_CACHE_MAX);
return collections;
apps/desktop/src/types/database.ts
+12 -1
View File
@@ -14,6 +14,7 @@ export type DatabaseType =
| "qdrant"
| "milvus"
| "weaviate"
| "chromadb"
| "doris"
| "starrocks"
| "manticoresearch"
@@ -544,7 +545,7 @@ export interface TreeNode {
hiddenChildren?: TreeNode[];
savedSqlId?: string;
savedSqlFolderId?: string;
meta?: ColumnInfo | IndexInfo | ForeignKeyInfo | TriggerInfo;
meta?: ColumnInfo | IndexInfo | ForeignKeyInfo | TriggerInfo | VectorCollectionMeta;
loadMore?: {
parentId: string;
offset: number;
@@ -682,3 +683,13 @@ export interface SavedSqlLibrary {
folders: SavedSqlFolder[];
files: SavedSqlFile[];
}
export interface VectorCollectionMeta {
dimension?: number;
}
export interface CollectionInfo {
name: string;
id: string;
dimension?: number;
}
crates/dbx-core/assets/database-drivers.manifest.json
+29
View File
@@ -349,6 +349,35 @@
"driverManagement": false
}
},
{
"dbType": "chromadb",
"label": "ChromaDB",
"runtimeMode": "native",
"mcpMode": "bridge",
"singleConnectionPool": false,
"metadataConnectionScoped": false,
"skipTcpProbe": false,
"defaultPort": 8000,
"supportLevel": "browse",
"capabilities": {
"queryExecution": true,
"metadataBrowse": true,
"objectBrowser": false,
"objectSource": false,
"schemaSearch": false,
"diagram": false,
"tableDataEdit": false,
"tableStructureEdit": false,
"tableImport": false,
"dataTransfer": false,
"sqlFileExecution": false,
"databaseCreate": false,
"fieldLineage": false,
"sqlExplain": false,
"userAdmin": false,
"driverManagement": false
}
},
{
"dbType": "milvus",
"label": "Milvus",
crates/dbx-core/src/connection.rs
+7 -2
View File
@@ -832,11 +832,12 @@ impl AppState {
db::elasticsearch_driver::test_connection(&mut client, connect_timeout).await?;
PoolKind::Elasticsearch(client)
}
DatabaseType::Qdrant | DatabaseType::Milvus | DatabaseType::Weaviate => {
DatabaseType::Qdrant | DatabaseType::Milvus | DatabaseType::Weaviate | DatabaseType::ChromaDb => {
let kind = match db_config.db_type {
DatabaseType::Qdrant => db::vector_driver::VectorDbKind::Qdrant,
DatabaseType::Milvus => db::vector_driver::VectorDbKind::Milvus,
DatabaseType::Weaviate => db::vector_driver::VectorDbKind::Weaviate,
DatabaseType::ChromaDb => db::vector_driver::VectorDbKind::ChromaDb,
_ => unreachable!(),
};
let client = db::vector_driver::VectorClient::new(
@@ -2140,7 +2141,11 @@ fn base_pool_key_for(
|| (include_elasticsearch_single_pool
&& matches!(
db_type,
DatabaseType::Elasticsearch | DatabaseType::Qdrant | DatabaseType::Milvus | DatabaseType::Weaviate
DatabaseType::Elasticsearch
| DatabaseType::Qdrant
| DatabaseType::Milvus
| DatabaseType::Weaviate
| DatabaseType::ChromaDb
));
is_single && (!database_capabilities::is_agent_type(db_type) || shares_database_pool_with_connection(db_type))
});
crates/dbx-core/src/db/vector_driver.rs
+190 -17
View File
@@ -27,11 +27,19 @@ const PATH_SEGMENT_ENCODE_SET: &AsciiSet = &CONTROLS
const QUERY_VALUE_ENCODE_SET: &AsciiSet = &PATH_SEGMENT_ENCODE_SET.add(b'&').add(b'=').add(b'+');
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct CollectionInfo {
pub name: String,
pub id: String,
pub dimension: Option<u32>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VectorDbKind {
Qdrant,
Milvus,
Weaviate,
ChromaDb,
}
impl VectorDbKind {
@@ -40,6 +48,7 @@ impl VectorDbKind {
VectorDbKind::Qdrant => "Qdrant",
VectorDbKind::Milvus => "Milvus",
VectorDbKind::Weaviate => "Weaviate",
VectorDbKind::ChromaDb => "ChromaDB",
}
}
}
@@ -57,6 +66,7 @@ enum VectorAuth {
Basic(String, String),
Bearer(String),
ApiKey(String),
ChromaToken(String),
}
impl VectorClient {
@@ -96,6 +106,7 @@ impl VectorClient {
Some(VectorAuth::Basic(user, pass)) => req.basic_auth(user, Some(pass)),
Some(VectorAuth::Bearer(token)) => req.bearer_auth(token),
Some(VectorAuth::ApiKey(token)) => req.header("api-key", token),
Some(VectorAuth::ChromaToken(token)) => req.header("x-chroma-token", token),
None => req,
}
}
@@ -114,6 +125,8 @@ fn vector_auth(kind: VectorDbKind, username: Option<&str>, password: Option<&str
VectorDbKind::Milvus => None,
VectorDbKind::Weaviate if !password.is_empty() => Some(VectorAuth::Bearer(password.to_string())),
VectorDbKind::Weaviate => None,
VectorDbKind::ChromaDb if !password.is_empty() => Some(VectorAuth::ChromaToken(password.to_string())),
VectorDbKind::ChromaDb => None,
}
}
@@ -123,11 +136,13 @@ pub async fn test_connection(client: &VectorClient, timeout: Duration) -> Result
VectorDbKind::Qdrant => "/collections",
VectorDbKind::Milvus => "/v2/vectordb/collections/list",
VectorDbKind::Weaviate => "/v1/meta",
VectorDbKind::ChromaDb => "/api/v2/heartbeat",
};
let request = match client.kind {
VectorDbKind::Qdrant => client.get(path),
VectorDbKind::Milvus => client.post(path).json(&serde_json::json!({ "dbName": "default" })),
VectorDbKind::Weaviate => client.get(path),
VectorDbKind::ChromaDb => client.get(path),
};
let resp = with_connection_timeout(label, timeout, async {
request.send().await.map_err(|e| format!("{label} connection failed: {}", format_reqwest_error(&e)))
@@ -136,39 +151,49 @@ pub async fn test_connection(client: &VectorClient, timeout: Duration) -> Result
ensure_success(label, resp).await.map(|_| ())
}
pub async fn list_collections(client: &VectorClient) -> Result<Vec<String>, String> {
pub async fn list_collections(client: &VectorClient) -> Result<Vec<CollectionInfo>, String> {
match client.kind {
VectorDbKind::Qdrant => list_qdrant_collections(client).await,
VectorDbKind::Milvus => list_milvus_collections(client).await,
VectorDbKind::Weaviate => list_weaviate_collections(client).await,
VectorDbKind::ChromaDb => list_chroma_collections(client).await,
}
}
async fn list_qdrant_collections(client: &VectorClient) -> Result<Vec<String>, String> {
async fn list_qdrant_collections(client: &VectorClient) -> Result<Vec<CollectionInfo>, String> {
let body = send_json(client.get("/collections"), "Qdrant").await?;
let mut names: Vec<String> = body
let mut infos: Vec<CollectionInfo> = body
.pointer("/result/collections")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(|item| item.get("name").and_then(Value::as_str).map(str::to_string))
.filter_map(|item| {
let name = item.get("name").and_then(Value::as_str)?;
Some(CollectionInfo { name: name.to_string(), id: name.to_string(), dimension: None })
})
.collect();
names.sort();
Ok(names)
infos.sort_by(|a, b| a.name.cmp(&b.name));
Ok(infos)
}
async fn list_milvus_collections(client: &VectorClient) -> Result<Vec<String>, String> {
async fn list_milvus_collections(client: &VectorClient) -> Result<Vec<CollectionInfo>, String> {
let body = send_json(
client.post("/v2/vectordb/collections/list").json(&serde_json::json!({ "dbName": "default" })),
"Milvus",
)
.await?;
let mut names: Vec<String> = match body.get("data") {
Some(Value::Array(items)) => items.iter().filter_map(collection_name_from_milvus_item).collect(),
let mut infos: Vec<CollectionInfo> = match body.get("data") {
Some(Value::Array(items)) => items
.iter()
.filter_map(|item| {
let name = collection_name_from_milvus_item(item)?;
Some(CollectionInfo { name: name.clone(), id: name, dimension: None })
})
.collect(),
_ => Vec::new(),
};
names.sort();
Ok(names)
infos.sort_by(|a, b| a.name.cmp(&b.name));
Ok(infos)
}
fn collection_name_from_milvus_item(item: &Value) -> Option<String> {
@@ -178,11 +203,56 @@ fn collection_name_from_milvus_item(item: &Value) -> Option<String> {
.or_else(|| item.get("name").and_then(Value::as_str).map(str::to_string))
}
async fn list_weaviate_collections(client: &VectorClient) -> Result<Vec<String>, String> {
async fn list_weaviate_collections(client: &VectorClient) -> Result<Vec<CollectionInfo>, String> {
let body = send_json(client.get("/v1/schema"), "Weaviate").await?;
let mut names = weaviate_collection_names_from_schema(&body);
names.sort();
Ok(names)
let mut infos: Vec<CollectionInfo> = weaviate_collection_names_from_schema(&body)
.into_iter()
.map(|name| CollectionInfo { name: name.clone(), id: name, dimension: None })
.collect();
infos.sort_by(|a, b| a.name.cmp(&b.name));
Ok(infos)
}
async fn list_chroma_collections(client: &VectorClient) -> Result<Vec<CollectionInfo>, String> {
let body =
send_json(client.get("/api/v2/tenants/default_tenant/databases/default_database/collections"), "ChromaDB")
.await?;
let mut infos: Vec<CollectionInfo> = body
.as_array()
.into_iter()
.flatten()
.filter_map(|item| {
let name = item.get("name").and_then(Value::as_str)?;
let id = item.get("id").and_then(Value::as_str)?;
let dimension = item.get("dimension").and_then(|v| v.as_u64()).map(|d| d as u32);
Some(CollectionInfo { name: name.to_string(), id: id.to_string(), dimension })
})
.collect();
infos.sort_by(|a, b| a.name.cmp(&b.name));
Ok(infos)
}
fn chroma_get_response_to_rows(body: &Value) -> Vec<Value> {
let ids = body.get("ids").and_then(Value::as_array).cloned().unwrap_or_default();
let docs = body.get("documents").and_then(Value::as_array).cloned().unwrap_or_default();
let metas = body.get("metadatas").and_then(Value::as_array).cloned().unwrap_or_default();
ids.into_iter()
.enumerate()
.map(|(i, id_val)| {
let mut row = serde_json::Map::new();
row.insert("id".to_string(), id_val);
if let Some(doc) = docs.get(i) {
row.insert("document".to_string(), doc.clone());
}
if let Some(Value::Object(meta_obj)) = metas.get(i) {
for (k, v) in meta_obj {
row.insert(k.clone(), v.clone());
}
}
Value::Object(row)
})
.collect()
}
fn weaviate_collection_names_from_schema(body: &Value) -> Vec<String> {
@@ -200,6 +270,45 @@ pub async fn find_documents(
skip: u64,
limit: i64,
) -> Result<crate::db::mongo_driver::MongoDocumentResult, String> {
if client.kind == VectorDbKind::ChromaDb {
let start = std::time::Instant::now();
let url = format!(
"{}/api/v2/tenants/default_tenant/databases/default_database/collections/{}/get",
client.base_url,
path_segment(collection),
);
let resp = client
.with_auth(client.http.post(&url))
.json(&serde_json::json!({
"limit": limit.max(1) as u64,
"offset": skip,
"include": ["documents", "metadatas"],
}))
.send()
.await
.map_err(|e| format!("ChromaDB request failed: {}", format_reqwest_error(&e)))?;
if !resp.status().is_success() {
let status = resp.status();
let body = resp.text().await.unwrap_or_default();
return Err(format!("ChromaDB error ({status}): {body}"));
}
let body: Value = resp.json().await.unwrap_or(Value::Null);
let rows = chroma_get_response_to_rows(&body);
let result = values_to_query_result(rows, start);
let documents = result
.rows
.into_iter()
.map(|row| {
let mut map = serde_json::Map::new();
for (idx, col) in result.columns.iter().enumerate() {
map.insert(col.clone(), row.get(idx).cloned().unwrap_or(Value::Null));
}
Value::Object(map)
})
.collect();
return Ok(crate::db::mongo_driver::MongoDocumentResult { documents, total: result.affected_rows });
}
let query = match client.kind {
VectorDbKind::Qdrant => format!(
"POST /collections/{}/points/scroll\n{}",
@@ -225,6 +334,7 @@ pub async fn find_documents(
VectorDbKind::Weaviate => {
format!("GET /v1/objects?class={}&limit={}&offset={}", query_value(collection), limit.max(1), skip)
}
VectorDbKind::ChromaDb => unreachable!("ChromaDB handled above"),
};
let result = execute_rest_query(client, &query).await?;
let documents = result
@@ -301,6 +411,12 @@ fn default_collection_query(client: &VectorClient, collection: &str) -> Result<r
"outputFields": ["*"],
}))),
VectorDbKind::Weaviate => Ok(client.get(&format!("/v1/objects?class={}&limit=100", query_value(collection)))),
VectorDbKind::ChromaDb => Ok(client
.post(&format!(
"/api/v2/tenants/default_tenant/databases/default_database/collections/{}/get",
path_segment(collection)
))
.json(&serde_json::json!({"limit": 100, "include": ["documents", "metadatas"]}))),
}
}
@@ -343,6 +459,11 @@ fn json_to_query_result(status: u16, body: Value, start: Instant) -> QueryResult
if let Some(Value::Array(items)) = body.pointer("/result/collections").cloned() {
return values_to_query_result(items, start);
}
if body.get("ids").and_then(Value::as_array).is_some() && body.get("documents").and_then(Value::as_array).is_some()
{
let rows = chroma_get_response_to_rows(&body);
return values_to_query_result(rows, start);
}
QueryResult {
columns: vec!["status".to_string(), "response".to_string()],
column_types: Vec::new(),
@@ -428,8 +549,8 @@ fn format_reqwest_error(err: &reqwest::Error) -> String {
#[cfg(test)]
mod tests {
use super::{
starts_with_http_method, values_to_query_result, vector_auth, weaviate_collection_names_from_schema,
VectorAuth, VectorDbKind,
chroma_get_response_to_rows, starts_with_http_method, values_to_query_result, vector_auth,
weaviate_collection_names_from_schema, CollectionInfo, VectorAuth, VectorDbKind,
};
use serde_json::json;
use std::time::Instant;
@@ -479,4 +600,56 @@ mod tests {
Some(VectorAuth::Bearer("token".to_string()))
);
}
#[test]
fn chroma_db_uses_x_chroma_token_header() {
assert_eq!(
vector_auth(VectorDbKind::ChromaDb, None, Some("my-key")),
Some(VectorAuth::ChromaToken("my-key".to_string()))
);
}
#[test]
fn chroma_db_no_auth_when_no_password() {
assert_eq!(vector_auth(VectorDbKind::ChromaDb, None, None), None);
}
#[test]
fn parses_chroma_collection_list() {
let body = json!([
{"id": "uuid-123", "name": "my_collection", "dimension": 384},
{"id": "uuid-456", "name": "other", "dimension": 768}
]);
let infos: Vec<CollectionInfo> = body
.as_array()
.unwrap()
.iter()
.filter_map(|item| {
let name = item.get("name").and_then(|v| v.as_str())?;
let id = item.get("id").and_then(|v| v.as_str())?;
Some(CollectionInfo { name: name.to_string(), id: id.to_string(), dimension: None })
})
.collect();
assert_eq!(infos.len(), 2);
assert_eq!(infos[0].name, "my_collection");
assert_eq!(infos[0].id, "uuid-123");
assert_eq!(infos[1].name, "other");
assert_eq!(infos[1].id, "uuid-456");
}
#[test]
fn converts_chroma_column_major_to_rows() {
let body = json!({
"ids": ["id1", "id2"],
"documents": ["hello world", "test doc"],
"metadatas": [{"source": "test"}, {"source": "demo"}]
});
let rows = chroma_get_response_to_rows(&body);
assert_eq!(rows.len(), 2);
assert_eq!(rows[0]["id"], json!("id1"));
assert_eq!(rows[0]["document"], json!("hello world"));
assert_eq!(rows[0]["source"], json!("test"));
assert_eq!(rows[1]["id"], json!("id2"));
assert_eq!(rows[1]["source"], json!("demo"));
}
}
crates/dbx-core/src/models/connection.rs
+12 -2
View File
@@ -263,6 +263,8 @@ pub enum DatabaseType {
Milvus,
#[serde(rename = "weaviate")]
Weaviate,
#[serde(rename = "chromadb")]
ChromaDb,
Doris,
#[serde(rename = "starrocks")]
StarRocks,
@@ -743,7 +745,11 @@ impl ConnectionConfig {
format!("mongodb://{host}:{port}{db_part}{suffix}")
}
DatabaseType::Oracle => format!("oracle://{host}:{port}{db_part}"),
DatabaseType::Elasticsearch | DatabaseType::Qdrant | DatabaseType::Milvus | DatabaseType::Weaviate => {
DatabaseType::Elasticsearch
| DatabaseType::Qdrant
| DatabaseType::Milvus
| DatabaseType::Weaviate
| DatabaseType::ChromaDb => {
let scheme = if self.ssl { "https" } else { "http" };
format!("{scheme}://{host}:{port}")
}
@@ -882,7 +888,11 @@ impl ConnectionConfig {
DatabaseType::Oracle => {
format!("oracle://{}:{}@{host}:{port}{db_part}", username, password)
}
DatabaseType::Elasticsearch | DatabaseType::Qdrant | DatabaseType::Milvus | DatabaseType::Weaviate => {
DatabaseType::Elasticsearch
| DatabaseType::Qdrant
| DatabaseType::Milvus
| DatabaseType::Weaviate
| DatabaseType::ChromaDb => {
let scheme = if self.ssl { "https" } else { "http" };
format!("{scheme}://{host}:{port}")
}
crates/dbx-core/src/mongo_ops.rs
+14 -5
View File
@@ -60,20 +60,29 @@ fn mongo_list_databases_unauthorized(error: &str) -> bool {
lower.contains("not authorized") && lower.contains("listdatabases")
}
use crate::db::vector_driver::CollectionInfo;
pub async fn mongo_list_collections_core(
state: &AppState,
connection_id: &str,
database: &str,
) -> Result<Vec<String>, String> {
) -> Result<Vec<CollectionInfo>, String> {
ensure_document_pool(state, connection_id).await?;
let connections = state.connections.read().await;
match connections.get(connection_id).ok_or("Not found")? {
PoolKind::MongoDb(client) => mongo_driver::list_collections(client, database).await.map(sort_names),
PoolKind::Elasticsearch(client) => elasticsearch_driver::list_indices(client).await.map(sort_names),
PoolKind::VectorDb(client) => vector_driver::list_collections(client).await.map(sort_names),
PoolKind::MongoDb(client) => {
let names = sort_names(mongo_driver::list_collections(client, database).await?);
Ok(names.into_iter().map(|n| CollectionInfo { name: n.clone(), id: n, dimension: None }).collect())
}
PoolKind::Elasticsearch(client) => {
let names = sort_names(elasticsearch_driver::list_indices(client).await?);
Ok(names.into_iter().map(|n| CollectionInfo { name: n.clone(), id: n, dimension: None }).collect())
}
PoolKind::VectorDb(client) => vector_driver::list_collections(client).await,
PoolKind::Agent(client) => {
let mut client = client.lock().await;
client.mongo_list_collections(database).await.map(sort_names)
let names = sort_names(client.mongo_list_collections(database).await?);
Ok(names.into_iter().map(|n| CollectionInfo { name: n.clone(), id: n, dimension: None }).collect())
}
_ => Err("Not a MongoDB/Elasticsearch/vector connection".to_string()),
}
crates/dbx-core/src/query.rs
+1
View File
@@ -734,6 +734,7 @@ fn should_discard_pool_after_query_timeout(db_type: Option<DatabaseType>) -> boo
| DatabaseType::Qdrant
| DatabaseType::Milvus
| DatabaseType::Weaviate
| DatabaseType::ChromaDb
| DatabaseType::InfluxDb
)
}
crates/dbx-core/src/query_execution_sql.rs
+1
View File
@@ -90,6 +90,7 @@ pub fn supports_sql_query(database_type: DatabaseType) -> bool {
| DatabaseType::Qdrant
| DatabaseType::Milvus
| DatabaseType::Weaviate
| DatabaseType::ChromaDb
| DatabaseType::InfluxDb
| DatabaseType::Neo4j
| DatabaseType::Etcd
crates/dbx-core/src/schema.rs
+1 -1
View File
@@ -1057,7 +1057,7 @@ async fn list_tables_once(
.map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)),
PoolKind::VectorDb(client) => db::vector_driver::list_collections(client)
.await
.map(|names| collection_names_to_tables(names, "COLLECTION"))
.map(|infos| collection_names_to_tables(infos.into_iter().map(|i| i.name).collect(), "COLLECTION"))
.map(|tables| filter_table_infos(tables, filter, limit, offset, object_types)),
_ => Ok(vec![]),
}
crates/dbx-core/src/schema/providers/native.rs
+1 -1
View File
@@ -59,7 +59,7 @@ pub(in crate::schema) async fn list_tables(
db::elasticsearch_driver::list_indices(client).await.map(|names| collection_names_to_tables(names, "INDEX"))
}
PoolKind::VectorDb(client) => {
db::vector_driver::list_collections(client).await.map(|names| collection_names_to_tables(names, "COLLECTION"))
db::vector_driver::list_collections(client).await.map(|infos| collection_names_to_tables(infos.into_iter().map(|i| i.name).collect(), "COLLECTION"))
}
_ => Ok(vec![]),
}
crates/dbx-web/src/routes/mongo.rs
+1 -1
View File
@@ -158,7 +158,7 @@ pub async fn list_databases(
pub async fn list_collections(
State(state): State<Arc<WebState>>,
Json(req): Json<MongoCollectionRequest>,
) -> Result<Json<Vec<String>>, AppError> {
) -> Result<Json<Vec<dbx_core::db::vector_driver::CollectionInfo>>, AppError> {
let result = dbx_core::mongo_ops::mongo_list_collections_core(&state.app, &req.connection_id, &req.database)
.await
.map_err(AppError)?;
packages/mcp-server/src/index.ts
+1 -1
View File
@@ -46,7 +46,7 @@ function formatQueryToolResult(result: QueryResult, title?: string) {
}
export const DBX_CONNECTION_TYPE_DESCRIPTION =
"Database type: postgres, mysql, sqlite, rqlite, redis, duckdb, clickhouse, sqlserver, mongodb, oracle, elasticsearch, etcd, doris, starrocks, manticoresearch, milvus, qdrant, weaviate, redshift, dameng, kingbase, highgo, vastbase, goldendb, databend, gaussdb, kwdb, yashandb, databricks, saphana, teradata, vertica, firebird, exasol, opengauss, oceanbase-oracle, questdb, gbase, h2, snowflake, trino, prestosql, hive, db2, informix, influxdb, iris, neo4j, cassandra, bigquery, kylin, sundb, tdengine, iotdb, xugu, zookeeper, jdbc, access, mq";
"Database type: postgres, mysql, sqlite, rqlite, redis, duckdb, clickhouse, sqlserver, mongodb, oracle, elasticsearch, etcd, doris, starrocks, manticoresearch, milvus, qdrant, weaviate, chromadb, redshift, dameng, kingbase, highgo, vastbase, goldendb, databend, gaussdb, kwdb, yashandb, databricks, saphana, teradata, vertica, firebird, exasol, opengauss, oceanbase-oracle, questdb, gbase, h2, snowflake, trino, prestosql, hive, db2, informix, influxdb, iris, neo4j, cassandra, bigquery, kylin, sundb, tdengine, iotdb, xugu, zookeeper, jdbc, access, mq";
const FILE_CAPABLE_CONNECTION_TYPES = new Set(["sqlite", "duckdb", "access", "h2"]);
interface McpScope {
packages/node-core/src/database.ts
+17 -2
View File
@@ -13,6 +13,14 @@ export interface TableInfo {
type: string;
}
export interface CollectionInfo {
name: string;
id?: string;
dimension?: number | null;
}
type CollectionListEntry = string | CollectionInfo;
export interface ColumnInfo {
name: string;
data_type: string;
@@ -408,6 +416,13 @@ interface BridgeColumnInfo {
character_maximum_length?: number | null;
}
export function collectionListToTableInfos(collections: CollectionListEntry[]): TableInfo[] {
return collections.map((collection) => ({
name: typeof collection === "string" ? collection : collection.name,
type: "COLLECTION",
}));
}
interface MongoDocumentResult {
documents: unknown[];
total: number;
@@ -646,12 +661,12 @@ export async function executeQuery(config: ConnectionConfig, sql: string, option
export async function listTables(config: ConnectionConfig, schema?: string): Promise<TableInfo[]> {
if (config.db_type === "mongodb") {
const collections = await bridgeDataRequest<string[]>("/data/mongo/list-collections", {
const collections = await bridgeDataRequest<CollectionListEntry[]>("/data/mongo/list-collections", {
connection_name: config.name,
database: config.database || "",
schema: schema || "",
});
return collections.map((name) => ({ name, type: "COLLECTION" }));
return collectionListToTableInfos(collections);
}
if (config.db_type === "sqlite" || config.db_type === "rqlite") {
const result = await query(config, `SELECT name, type FROM sqlite_master WHERE type IN ('table', 'view') AND name NOT LIKE 'sqlite_%' ORDER BY name`);
packages/node-core/src/diagnostics.ts
+1
View File
@@ -23,6 +23,7 @@ export const BRIDGE_REQUIRED_TYPES = [
"qdrant",
"milvus",
"weaviate",
"chromadb",
"etcd",
"dameng",
"kingbase",
packages/node-core/src/web-backend.ts
+3 -3
View File
@@ -1,6 +1,6 @@
import type { ConnectionConfig } from "./connections.js";
import type { TableInfo, ColumnInfo, QueryOptions, QueryResult } from "./database.js";
import { evaluateMongoAggregateSafety, evaluateMongoWriteSafety, inferMongoColumns, mongoDocumentsToQueryResult, parseMongoAggregateCommand, parseMongoCountDocumentsCommand, parseMongoFindCommand, parseMongoGetIndexesCommand, parseMongoWriteCommand, type MongoWriteCommand } from "./database.js";
import { collectionListToTableInfos, evaluateMongoAggregateSafety, evaluateMongoWriteSafety, inferMongoColumns, mongoDocumentsToQueryResult, parseMongoAggregateCommand, parseMongoCountDocumentsCommand, parseMongoFindCommand, parseMongoGetIndexesCommand, parseMongoWriteCommand, type CollectionInfo, type MongoWriteCommand } from "./database.js";
import { sqlSafetyFromEnv } from "./sql-safety.js";
const baseUrl = process.env.DBX_WEB_URL!.replace(/\/+$/, "");
@@ -93,8 +93,8 @@ export async function listTables(config: ConnectionConfig, schema?: string): Pro
method: "POST",
body: JSON.stringify({ connectionId: config.id, database: config.database || "" }),
});
const collections = (await res.json()) as string[];
return collections.map((name) => ({ name, type: "COLLECTION" }));
const collections = (await res.json()) as Array<string | CollectionInfo>;
return collectionListToTableInfos(collections);
}
const params = new URLSearchParams({
connection_id: config.id,
packages/node-core/tests/collections.test.ts
+23
View File
@@ -0,0 +1,23 @@
import assert from "node:assert/strict";
import { test } from "vitest";
import { collectionListToTableInfos } from "../src/database.js";
test("maps legacy collection name responses to table infos", () => {
assert.deepEqual(collectionListToTableInfos(["projects", "users"]), [
{ name: "projects", type: "COLLECTION" },
{ name: "users", type: "COLLECTION" },
]);
});
test("maps collection info responses to table infos", () => {
assert.deepEqual(
collectionListToTableInfos([
{ name: "projects", id: "projects", dimension: null },
{ name: "embeddings", id: "uuid-123", dimension: 384 },
]),
[
{ name: "projects", type: "COLLECTION" },
{ name: "embeddings", type: "COLLECTION" },
],
);
});
src-tauri/src/commands/connection.rs
+4 -2
View File
@@ -713,11 +713,12 @@ pub async fn test_connection(state: State<'_, Arc<AppState>>, config: Connection
.await
.map(|_| "Connection successful".to_string())
}
DatabaseType::Qdrant | DatabaseType::Milvus | DatabaseType::Weaviate => {
DatabaseType::Qdrant | DatabaseType::Milvus | DatabaseType::Weaviate | DatabaseType::ChromaDb => {
let kind = match config.db_type {
DatabaseType::Qdrant => db::vector_driver::VectorDbKind::Qdrant,
DatabaseType::Milvus => db::vector_driver::VectorDbKind::Milvus,
DatabaseType::Weaviate => db::vector_driver::VectorDbKind::Weaviate,
DatabaseType::ChromaDb => db::vector_driver::VectorDbKind::ChromaDb,
_ => unreachable!(),
};
let client = db::vector_driver::VectorClient::new(
@@ -999,11 +1000,12 @@ pub async fn connect_db(state: State<'_, Arc<AppState>>, config: ConnectionConfi
db::elasticsearch_driver::test_connection(&mut client, connect_timeout).await?;
PoolKind::Elasticsearch(client)
}
DatabaseType::Qdrant | DatabaseType::Milvus | DatabaseType::Weaviate => {
DatabaseType::Qdrant | DatabaseType::Milvus | DatabaseType::Weaviate | DatabaseType::ChromaDb => {
let kind = match db_config.db_type {
DatabaseType::Qdrant => db::vector_driver::VectorDbKind::Qdrant,
DatabaseType::Milvus => db::vector_driver::VectorDbKind::Milvus,
DatabaseType::Weaviate => db::vector_driver::VectorDbKind::Weaviate,
DatabaseType::ChromaDb => db::vector_driver::VectorDbKind::ChromaDb,
_ => unreachable!(),
};
let client = db::vector_driver::VectorClient::new(
src-tauri/src/commands/mongo_cmd.rs
+1 -1
View File
@@ -36,7 +36,7 @@ pub async fn mongo_list_collections(
state: State<'_, Arc<AppState>>,
connection_id: String,
database: String,
) -> Result<Vec<String>, String> {
) -> Result<Vec<dbx_core::db::vector_driver::CollectionInfo>, String> {
dbx_core::mongo_ops::mongo_list_collections_core(&state, &connection_id, &database).await
}