Skip to content

chore(poc): update PoC to use new WASM interface (SQL-3170) #8119

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
nicholascioli wants to merge 1 commit into
base: schema-builder-poc
Choose a base branch
from
Open

chore(poc): update PoC to use new WASM interface (SQL-3170) #8119

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
122 changes: 57 additions & 65 deletions packages/compass-data-modeling/src/services/compass-data-service-adapter.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,57 +1,55 @@
import type { DataService } from 'mongodb-data-service';
import { EJSON, type Document } from 'bson';

/**
* BSON Document type - a generic key-value object.
* This mirrors the BsonDocument type from schema-builder-library to avoid ESM/CJS issues.
* When crossing the WASM boundary, BSON types are represented in EJSON format
* (e.g., ObjectId as { "$oid": "..." }).
*/
export type BsonDocument = Record<string, unknown>;
import { AggregationCursor, FindCursor } from 'mongodb';
import type { DataService } from 'mongodb-data-service';
import type {
AggregateOptions,
BsonDocument,
CollectionInfo,
SqlCursor,
SqlDataService,
} from 'schema-builder-library';

/**
* Convert native BSON documents to EJSON format for WASM.
* This ensures BSON types like ObjectId, Date, etc. are properly serialized
* as EJSON objects that can be deserialized by Rust's bson crate.
*/
function toEJSON(docs: Document[]): BsonDocument[] {
function toEJSON(doc: Document): BsonDocument {
// EJSON.serialize converts BSON types to their EJSON representation
// e.g., ObjectId("...") -> { "$oid": "..." }
return docs.map((doc) => EJSON.serialize(doc) as BsonDocument);
return EJSON.serialize(doc, { relaxed: false }) as BsonDocument;
}

/**
* Convert EJSON format from WASM back to native BSON documents.
* This ensures EJSON objects like { "$oid": "..." } are converted back
* to native BSON types like ObjectId for the MongoDB driver.
*/
function fromEJSON(docs: BsonDocument[]): Document[] {
function fromEJSON(doc: BsonDocument): Document {
// EJSON.deserialize converts EJSON representation back to BSON types
// e.g., { "$oid": "..." } -> ObjectId("...")
return docs.map((doc) => EJSON.deserialize(doc as Document) as Document);
return EJSON.deserialize(doc as Document);
}

/**
* Collection info type - mirrors CollectionInfo from schema-builder-library.
* A wrapper around the node driver cursor that automatically converts the
* native EJSON results to WASM-friendly ones.
*/
export interface WasmCollectionInfo {
name: string;
type: string;
options?: {
viewOn?: string;
pipeline?: BsonDocument[];
};
}
class Cursor implements SqlCursor {
cursor: AggregationCursor | FindCursor;

/**
* DataService interface - mirrors DataService from schema-builder-library.
* This interface is what the WASM module expects.
*/
export interface WasmDataService {
listDatabases(): Promise<string[]>;
listCollections(dbName: string): Promise<WasmCollectionInfo[]>;
aggregate(ns: string, pipeline: BsonDocument[]): Promise<BsonDocument[]>;
find(ns: string, filter: BsonDocument): Promise<BsonDocument[]>;
constructor(cursor: AggregationCursor | FindCursor) {
this.cursor = cursor;
}

async next(): Promise<BsonDocument | null> {
let next = (await this.cursor.next()) as Document;
if (!next) {
return null;
}

return toEJSON(next);
}
}

/**
Expand All @@ -61,7 +59,7 @@ export interface WasmDataService {
* This allows the Rust-based schema builder to use Compass's existing database connection
* without needing direct access to the MongoDB driver.
*/
export class CompassDataServiceAdapter implements WasmDataService {
export class CompassDataServiceAdapter implements SqlDataService {
private dataService: DataService;
private abortSignal?: AbortSignal;

Expand All @@ -77,36 +75,29 @@ export class CompassDataServiceAdapter implements WasmDataService {
const databases = await this.dataService.listDatabases({
nameOnly: true,
});

return databases.map((db) => db.name);
}

/**
* List all collections in a database, including views.
* Returns collection info with name, type, and options (for views).
*/
async listCollections(dbName: string): Promise<WasmCollectionInfo[]> {
async listCollections(dbName: string): Promise<CollectionInfo[]> {
const collections = await this.dataService.listCollections(dbName);

return collections.map((coll) => {
const collInfo: WasmCollectionInfo = {
const collInfo: CollectionInfo = {
name: coll.name,
type: coll.type ?? 'collection',
};

// For views, include viewOn and pipeline in options
if (coll.type === 'view') {
const collectionInfo = coll as unknown as {
options?: {
viewOn?: string;
pipeline?: Document[];
};
collInfo.options = {
viewOn: coll.view_on ?? undefined,
pipeline: coll.pipeline?.map(toEJSON),
};
if (collectionInfo.options) {
collInfo.options = {
viewOn: collectionInfo.options.viewOn,
pipeline: collectionInfo.options.pipeline,
};
}
}

return collInfo;
Expand All @@ -121,23 +112,24 @@ export class CompassDataServiceAdapter implements WasmDataService {
* Results are converted back to EJSON format for WASM.
*/
async aggregate(
ns: string,
pipeline: BsonDocument[]
): Promise<BsonDocument[]> {
db: string,
collection: string,
pipeline: BsonDocument[],
options: Partial<AggregateOptions>
): Promise<SqlCursor> {
// Convert EJSON pipeline from WASM to native BSON for DataService
const nativePipeline = fromEJSON(pipeline);
const nativePipeline = pipeline.map(fromEJSON);

const result = await this.dataService.aggregate(
ns,
const result = this.dataService.aggregateCursor(
`${db}.${collection}`,
nativePipeline,
{},
{
abortSignal: this.abortSignal,
signal: this.abortSignal,
hint: options.keyHint,
}
);

// Convert native BSON results to EJSON for WASM
return toEJSON(result);
return new Cursor(result);
}

/**
Expand All @@ -147,20 +139,20 @@ export class CompassDataServiceAdapter implements WasmDataService {
* Filter comes from WASM in EJSON format and must be converted to native BSON.
* Results are converted back to EJSON format for WASM.
*/
async find(ns: string, filter: BsonDocument): Promise<BsonDocument[]> {
async find(
db: string,
collection: string,
filter: BsonDocument
): Promise<SqlCursor> {
// Convert EJSON filter from WASM to native BSON for DataService
const nativeFilter = EJSON.deserialize(filter as Document) as Document;
const nativeFilter = fromEJSON(filter);

const documents = await this.dataService.find(
ns,
const documents = this.dataService.findCursor(
`${db}.${collection}`,
nativeFilter,
{},
{
abortSignal: this.abortSignal,
}
{}
);

// Convert native BSON results to EJSON for WASM
return toEJSON(documents);
return new Cursor(documents);
}
}
122 changes: 30 additions & 92 deletions packages/compass-data-modeling/src/store/analysis-process.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ import {
getCurrentDiagramFromState,
selectCurrentModel,
} from './diagram';
import { UUID } from 'bson';
import { EJSON, UUID } from 'bson';
import {
type Relationship,
type StaticModel,
Expand All @@ -21,7 +21,10 @@ import { inferForeignToLocalRelationshipsForCollection } from './relationships';
import { mongoLogId } from '@mongodb-js/compass-logging/provider';
import { extractFieldsFromFieldData } from '../utils/schema';
import { isEqual } from 'lodash';
import type { WasmDataService } from '../services/compass-data-service-adapter';
import type {
process_collection,
SqlDataService,
} from 'schema-builder-library';
import { CompassDataServiceAdapter } from '../services/compass-data-service-adapter';

// WASM module lazy initialization
Expand All @@ -37,11 +40,7 @@ type WasmBuilderOptionsInstance = {
};

type WasmModule = {
buildSchema: (
dataService: WasmDataService,
options: WasmBuilderOptionsInstance
) => Promise<WasmSchemaResult[]>;
WasmBuilderOptions: new () => WasmBuilderOptionsInstance;
process_collection: typeof process_collection;
};

let wasmModulePromise: Promise<WasmModule> | null = null;
Expand All @@ -56,36 +55,6 @@ function getWasmModule(): Promise<WasmModule> {
return wasmModulePromise;
}

/**
* Type representing the serialized SchemaResult from the WASM module.
* The Rust enum SchemaResult serializes with the variant name as a key.
*/
export type WasmSchemaResult =
| { NamespaceOnly: WasmNamespaceInfo }
| { InitialSchema: WasmNamespaceInfoWithSchema }
| { FullSchema: WasmNamespaceInfoWithSchema };

type WasmNamespaceInfo = {
db_name: string;
coll_or_view_name: string;
namespace_type: 'Collection' | 'View';
};

type WasmNamespaceInfoWithSchema = {
namespace_info: WasmNamespaceInfo;
namespace_schema: MongoDBJSONSchema;
};

/**
* Options for building schemas.
*/
export interface SchemaBuilderOptions {
/** Data service for accessing MongoDB */
dataService: WasmDataService;
/** List of namespaces to include (format: "database.collection") */
includeList: string[];
}

/**
* Service interface for schema building.
* This allows the WASM-based schema builder to be mocked in tests.
Expand All @@ -96,28 +65,33 @@ export interface SchemaBuilderService {
* @param options - Schema builder options
* @returns Array of schema results
*/
buildSchemas(options: SchemaBuilderOptions): Promise<WasmSchemaResult[]>;
process_collection: (
service: SqlDataService,
ns: string
) => Promise<MongoDBJSONSchema | null>;
}

/**
* Default schema builder service using the WASM module.
*/
export const defaultSchemaBuilderService: SchemaBuilderService = {
async buildSchemas(
options: SchemaBuilderOptions
): Promise<WasmSchemaResult[]> {
async process_collection(service: SqlDataService, ns: string) {
try {
const wasmModule = await getWasmModule();
const { buildSchema, WasmBuilderOptions } = wasmModule;
const { process_collection } = wasmModule;

const wasmOptions = new WasmBuilderOptions();
wasmOptions.setIncludeList(options.includeList);
const [db, ...collParts] = ns.split('.');
const collection = collParts.join('.');

return await buildSchema(options.dataService, wasmOptions);
let result = await process_collection(service, db, collection);

// Note: the schema is returned as EJSON, which fails the serialization
// schema for the data provider, so we serialize it into regular JSON here
return EJSON.serialize(result);
} catch (err) {
// eslint-disable-next-line no-console
console.error('Failed to load WASM module', err);
return [];
return null;
}
},
};
Expand All @@ -127,27 +101,12 @@ export const defaultSchemaBuilderService: SchemaBuilderService = {
* Returns empty schemas for all requested collections.
*/
export const mockSchemaBuilderService: SchemaBuilderService = {
buildSchemas(options: SchemaBuilderOptions): Promise<WasmSchemaResult[]> {
// Return FullSchema results with empty schemas for each namespace
return Promise.resolve(
options.includeList.map((ns) => {
const [database, ...collParts] = ns.split('.');
const collection = collParts.join('.');
return {
FullSchema: {
namespace_info: {
db_name: database,
coll_or_view_name: collection,
namespace_type: 'Collection' as const,
},
namespace_schema: {
bsonType: 'object' as const,
properties: {},
},
},
};
})
);
async process_collection(service: SqlDataService, ns: string) {
// Returns an empty schema for any provided namespace
return {
bsonType: 'object' as const,
properties: {},
};
},
};

Expand Down Expand Up @@ -597,22 +556,6 @@ export function redoAnalysis(
};
}

/**
* Extract schema from a WASM SchemaResult.
* Returns the namespace info and schema if it's a FullSchema or InitialSchema variant.
*/
function extractSchemaFromResult(
result: WasmSchemaResult
): WasmNamespaceInfoWithSchema | null {
if ('FullSchema' in result) {
return result.FullSchema;
}
if ('InitialSchema' in result) {
return result.InitialSchema;
}
return null;
}

export function analyzeCollections({
name,
connectionId,
Expand Down Expand Up @@ -671,18 +614,13 @@ export function analyzeCollections({

// Step 1: Build schemas using the schema builder service
const adapter = new CompassDataServiceAdapter(dataService, abortSignal);
const schemaResults: WasmSchemaResult[] = await schemaBuilder.buildSchemas({
dataService: adapter,
includeList: selectedCollections.map((coll) => `${database}.${coll}`),
});

// Create a map of namespace -> schema for quick lookup
const schemaMap = new Map<string, MongoDBJSONSchema>();
for (const result of schemaResults) {
const schemaInfo = extractSchemaFromResult(result);
if (schemaInfo) {
const ns = `${schemaInfo.namespace_info.db_name}.${schemaInfo.namespace_info.coll_or_view_name}`;
schemaMap.set(ns, schemaInfo.namespace_schema);
for (const ns of namespaces) {
const schema = await schemaBuilder.process_collection(adapter, ns);
if (schema) {
schemaMap.set(ns, schema);
}
}

Expand Down