feat: link sample id with table id

packages/dbml-core/__tests__/examples/model_structure/records.spec.ts

@@ -0,0 +1,103 @@
+import Parser from '../../../src/parse/Parser';
+import { test, beforeAll, describe, expect } from 'vitest';
+
+const DBML = `
+  Table users {
+    id integer [pk]
+    name varchar
+  }
+
+  Table posts {
+    id integer [pk]
+    content text
+  }
+
+  Records users(id, name) {
+    1, "Alice"
+    2, "Bob"
+  }
+
+  Records posts(id, content) {
+    1, "Hello World"
+  }
+`;
+
+describe('@dbml/core - model_structure - records', () => {
+  let database: ReturnType<Parser['parse']>;
+
+  beforeAll(() => {
+    database = new Parser().parse(DBML, 'dbmlv2');
+  });
+
+  test('database.records contains all records', () => {
+    expect(database.records.length).toBe(2);
+  });
+
+  test('table.records is populated after linking', () => {
+    const usersTable = database.schemas[0].tables.find((t) => t.name === 'users');
+    const postsTable = database.schemas[0].tables.find((t) => t.name === 'posts');
+
+    expect(usersTable!.records.length).toBe(1);
+    expect(postsTable!.records.length).toBe(1);
+  });
+
+  test('table.records holds the same object references as database.records', () => {
+    const usersTable = database.schemas[0].tables.find((t) => t.name === 'users');
+
+    expect(database.records).toContain(usersTable!.records[0]);
+  });
+
+  test('linked record has correct tableName and values', () => {
+    const usersTable = database.schemas[0].tables.find((t) => t.name === 'users');
+    const record = usersTable!.records[0];
+
+    expect(record.tableName).toBe('users');
+    expect(record.values.length).toBe(2);
+  });
+
+  test('linked record has tableId set to the table id', () => {
+    const usersTable = database.schemas[0].tables.find((t) => t.name === 'users');
+    expect(usersTable!.records[0].tableId).toBe(usersTable!.id);
+  });
+
+  test('table without records has empty records array', () => {
+    const db = new Parser().parse(`
+      Table empty_table {
+        id integer [pk]
+      }
+    `, 'dbmlv2');
+
+    expect(db.schemas[0].tables[0].records).toEqual([]);
+  });
+
+  describe('normalized model', () => {
+    test('table has recordIds in normalized model', () => {
+      const normalizedModel = database.normalize();
+      const usersTable = database.schemas[0].tables.find((t) => t.name === 'users');
+      const postsTable = database.schemas[0].tables.find((t) => t.name === 'posts');
+
+      expect(normalizedModel.tables[usersTable!.id].recordIds).toEqual([usersTable!.records[0].id]);
+      expect(normalizedModel.tables[postsTable!.id].recordIds).toEqual([postsTable!.records[0].id]);
+    });
+
+    test('normalized records contain tableId', () => {
+      const normalizedModel = database.normalize();
+      const usersTable = database.schemas[0].tables.find((t) => t.name === 'users');
+      const recordId = usersTable!.records[0].id;
+
+      expect(normalizedModel.records[recordId].tableId).toBe(usersTable!.id);
+    });
+
+    test('table without records has empty recordIds in normalized model', () => {
+      const db = new Parser().parse(`
+        Table empty_table {
+          id integer [pk]
+        }
+      `, 'dbmlv2');
+
+      const normalizedModel = db.normalize();
+      const table = db.schemas[0].tables[0];
+      expect(normalizedModel.tables[table.id].recordIds).toEqual([]);
+    });
+  });
+});

packages/dbml-core/src/model_structure/check.js

@@ -1,20 +1,31 @@
+/// <reference path="../../types/model_structure/check.d.ts" />
+// @ts-check
 import Element from './element';
 
+/**
+  * @type  Check
+  */
 class Check extends Element {
   constructor ({
     token, name, expression, table, column = null, injectedPartial = null,
   } = {}) {
     super(token);
     this.name = name;
+    /** @type {string} */
     this.expression = expression;
+    /** @type {import('../../types/model_structure/table').default} */
     this.table = table;
+    /** @type {import('../../types/model_structure/field').default} */
     this.column = column;
+    /** @type {import('../../types/model_structure/tablePartial').default} */
     this.injectedPartial = injectedPartial;
+    /** @type {import('../../types/model_structure/dbState').default} */
     this.dbState = this.table.dbState;
     this.generateId();
   }
 
   generateId () {
+    /** @type {number} */
     this.id = this.dbState.generateId('checkId');
   }
 
@@ -39,6 +50,9 @@ class Check extends Element {
     };
   }
 
+  /**
+   * @param {import('../../types/model_structure/database').NormalizedDatabase} model
+   */
   normalize (model) {
     model.checks[this.id] = {
       id: this.id,

packages/dbml-core/src/model_structure/database.js

@@ -13,6 +13,9 @@ import DbState from './dbState';
 import TablePartial from './tablePartial';
 
 class Database extends Element {
+  /**
+    * @param {import('../../types/model_structure/database').RawDatabase} param0
+    */
   constructor ({
     schemas = [],
     tables = [],
@@ -29,6 +32,7 @@ class Database extends Element {
     this.dbState = new DbState();
     this.generateId();
     this.hasDefaultSchema = false;
+    /** @type {import('../../types/model_structure/schema').default[]} */
     this.schemas = [];
     this.notes = [];
     this.note = project.note ? get(project, 'note.value', project.note) : null;
@@ -51,6 +55,7 @@ class Database extends Element {
     this.processSchemas(schemas);
     this.processSchemaElements(enums, ENUM);
     this.processSchemaElements(tables, TABLE);
+    this.linkRecordsToTables();
     this.processSchemaElements(notes, NOTE);
     this.processSchemaElements(refs, REF);
     this.processSchemaElements(tableGroups, TABLE_GROUP);
@@ -157,6 +162,28 @@ class Database extends Element {
     });
   }
 
+  linkRecordsToTables () {
+    // Build a map of [schemaName][tableName] -> table for O(1) lookup
+    const tableMap = {};
+    this.schemas.forEach((schema) => {
+      tableMap[schema.name] = {};
+      schema.tables.forEach((table) => {
+        tableMap[schema.name][table.name] = table;
+      });
+    });
+
+    // Link records to tables using the map
+    this.records.forEach((record) => {
+      // Fallback to 'public' if schemaName is null, undefined
+      const schemaName = record.schemaName ?? DEFAULT_SCHEMA_NAME;
+      const table = tableMap[schemaName]?.[record.tableName];
+      if (!table) return;
+
+      record.tableId = table.id;
+      table.records.push(record);
+    });
+  }
+
   findOrCreateSchema (schemaName) {
     let schema = this.schemas.find((s) => s.name === schemaName || s.alias === schemaName);
     // create new schema if schema not found

packages/dbml-core/src/model_structure/dbState.js

@@ -1,22 +1,41 @@
 export default class DbState {
   constructor () {
+    /** @type {number} */
     this.dbId = 1;
+    /** @type {number} */
     this.schemaId = 1;
+    /** @type {number} */
     this.enumId = 1;
+    /** @type {number} */
     this.tableGroupId = 1;
+    /** @type {number} */
     this.refId = 1;
+    /** @type {number} */
     this.tableId = 1;
+    /** @type {number} */
     this.noteId = 1;
+    /** @type {number} */
     this.enumValueId = 1;
+    /** @type {number} */
     this.endpointId = 1;
+    /** @type {number} */
     this.indexId = 1;
+    /** @type {number} */
     this.checkId = 1;
+    /** @type {number} */
     this.fieldId = 1;
+    /** @type {number} */
     this.indexColumnId = 1;
+    /** @type {number} */
     this.recordId = 1;
+    /** @type {number} */
     this.tablePartialId = 1;
   }
 
+  /**
+   * @param {string} el
+   * @returns {number}
+   */
   generateId (el) {
     const id = this[el];
     this[el] += 1;

packages/dbml-core/src/model_structure/element.js

@@ -1,20 +1,34 @@
 class ElementError extends Error {
+  /**
+   * @param {string} message
+   * @param {import('../../types/model_structure/element').Token} location
+   */
   constructor (message, location = { start: { line: 1, column: 1 } }) {
     super(message);
+    /** @type {import('../../types/model_structure/element').Token} */
     this.location = location;
+    /** @type {string} */
     this.error = 'error';
   }
 }
 
 class Element {
   constructor (token) {
+    /** @type {import('../../types/model_structure/element').Token} */
     this.token = token;
   }
 
+  /**
+   * @param {any} selection
+   */
   bind (selection) {
+    /** @type {any} */
     this.selection = selection;
   }
 
+  /**
+   * @param {string} message
+   */
   error (message) {
     throw new ElementError(message, this.token);
   }

packages/dbml-core/src/model_structure/endpoint.js

@@ -3,17 +3,27 @@ import { DEFAULT_SCHEMA_NAME } from './config';
 import { shouldPrintSchema, shouldPrintSchemaName } from './utils';
 
 class Endpoint extends Element {
+  /**
+   * @param {{ tableName: string, schemaName: string, fieldNames: string[], relation: any, token: import('../../types/model_structure/element').Token, ref: import('../../types/model_structure/ref').default }} param0
+   */
   constructor ({
     tableName, schemaName, fieldNames, relation, token, ref,
   }) {
     super(token);
+    /** @type {any} */
     this.relation = relation;
 
+    /** @type {string} */
     this.schemaName = schemaName;
+    /** @type {string} */
     this.tableName = tableName;
+    /** @type {string[]} */
     this.fieldNames = fieldNames;
+    /** @type {import('../../types/model_structure/field').default[]} */
     this.fields = [];
+    /** @type {import('../../types/model_structure/ref').default} */
     this.ref = ref;
+    /** @type {import('../../types/model_structure/dbState').default} */
     this.dbState = this.ref.dbState;
     this.generateId();
     // Use name of schema,table and field object
@@ -30,14 +40,23 @@ class Endpoint extends Element {
   }
 
   generateId () {
+    /** @type {number} */
     this.id = this.dbState.generateId('endpointId');
   }
 
+  /**
+   * @param {import('../../types/model_structure/endpoint').default} endpoint
+   * @returns {boolean}
+   */
   equals (endpoint) {
     if (this.fields.length !== endpoint.fields.length) return false;
     return this.compareFields(endpoint);
   }
 
+  /**
+   * @param {import('../../types/model_structure/endpoint').default} endpoint
+   * @returns {boolean}
+   */
   compareFields (endpoint) {
     const sortedThisFieldIds = this.fields.map((field) => field.id).sort();
     const sortedEndpointFieldIds = endpoint.fields.map((field) => field.id).sort();
@@ -69,6 +88,10 @@ class Endpoint extends Element {
     };
   }
 
+  /**
+   * @param {string[]} fieldNames
+   * @param {import('../../types/model_structure/table').default} table
+   */
   setFields (fieldNames, table) {
     let newFieldNames = (fieldNames && fieldNames.length) ? [...fieldNames] : [];
     if (!newFieldNames.length) {
@@ -96,6 +119,9 @@ class Endpoint extends Element {
     });
   }
 
+  /**
+   * @param {import('../../types/model_structure/database').NormalizedDatabase} model
+   */
   normalize (model) {
     model.endpoints[this.id] = {
       id: this.id,