feat: optionally include leading comments in statement text

README.md

@@ -152,6 +152,38 @@ console.log(statements);
 1. `options (object)`: allow to set different configurations
     1. `strict (bool)`: allow disable strict mode which will ignore unknown types *(default=true)*
     2. `dialect (string)`: Specify your database dialect, values: `generic`, `mysql`, `oracle`, `psql`, `sqlite`, `mssql`, `bigquery` and `dynamodb`. *(default=generic)*
+    3. `includeLeadingComments (bool)`: Include leading comments in statement `text` and `start`, starting at the first comment unless `includeLeadingWhitespace` is also enabled. *(default=false)*
+    4. `includeLeadingWhitespace (bool)`: Include whitespace immediately before the statement, or before its leading comments when `includeLeadingComments` is enabled. *(default=false)*
+
+### Leading comments and whitespace
+
+Both options default to `false`, so existing callers still receive text and
+offsets that start at the first keyword. Neither option changes `end`.
+
+For `  -- header\n  SELECT 1;`:
+
+| Options | Returned `text` starts at |
+| --- | --- |
+| Neither | `SELECT` |
+| `includeLeadingComments` | `-- header` |
+| `includeLeadingWhitespace` | The two spaces after the comment line |
+| Both | The two spaces before `-- header` |
+
+Comment-only mode keeps whitespace between comments and the keyword because
+`text` remains a contiguous source substring. Whitespace-only mode does not cross
+or include a comment. The tokenizer includes a line comment's terminating newline
+in that comment token, so whitespace-only mode does not retain that newline.
+Semicolon delimiters are excluded in every mode.
+
+To preserve a module's header comments and their leading whitespace:
+
+```js
+identify('  -- header\nCREATE VIEW v AS SELECT 1;', {
+  dialect: 'mssql',
+  includeLeadingComments: true,
+  includeLeadingWhitespace: true,
+});
+```
 
 ### DynamoDB (PartiQL)
 

src/defines.ts

@@ -136,6 +136,8 @@ export interface TableReference {
 
 export interface IdentifyOptions {
   strict?: boolean;
+  includeLeadingComments?: boolean;
+  includeLeadingWhitespace?: boolean;
   dialect?: Dialect;
   identifyTables?: boolean;
   identifyColumns?: boolean;

src/index.ts

@@ -24,7 +24,7 @@ export function identify(query: string, options: IdentifyOptions = {}): Identify
   // Default parameter types for each dialect
   const paramTypes = options.paramTypes || defaultParamTypesFor(dialect);
 
-  const result = parse(
+  const parsed = parse(
     query,
     isStrict,
     dialect,
@@ -34,11 +34,43 @@ export function identify(query: string, options: IdentifyOptions = {}): Identify
   );
   const sort = dialect === 'psql' && !options.paramTypes;
 
-  return result.body.map((statement) => {
+  // Use token boundaries so comment markers inside strings cannot become headers.
+  let tokenIndex = 0;
+  let previousEnd = -1;
+  return parsed.body.map((statement) => {
+    let start = statement.start;
+    let leadingStart: number | undefined;
+    let firstComment: number | undefined;
+    while (
+      (options.includeLeadingComments === true || options.includeLeadingWhitespace === true) &&
+      tokenIndex < parsed.tokens.length &&
+      parsed.tokens[tokenIndex].start < statement.start
+    ) {
+      const token = parsed.tokens[tokenIndex++];
+      if (token.end <= previousEnd) continue;
+      if (token.type === 'whitespace') {
+        if (leadingStart === undefined) leadingStart = token.start;
+      } else if (
+        options.includeLeadingComments === true &&
+        (token.type === 'comment-inline' || token.type === 'comment-block')
+      ) {
+        if (leadingStart === undefined) leadingStart = token.start;
+        if (firstComment === undefined) firstComment = token.start;
+      } else {
+        leadingStart = undefined;
+        firstComment = undefined;
+      }
+    }
+    if (options.includeLeadingWhitespace === true && leadingStart !== undefined) {
+      start = leadingStart;
+    } else if (options.includeLeadingComments === true && firstComment !== undefined) {
+      start = firstComment;
+    }
+    previousEnd = statement.end;
     const result: IdentifyResult = {
-      start: statement.start,
+      start,
       end: statement.end,
-      text: query.substring(statement.start, statement.end + 1),
+      text: query.substring(start, statement.end + 1),
       type: statement.type,
       executionType: statement.executionType,
       // we want to sort the postgres params: $1 $2 $3, regardless of the order they appear

test/identifier/single-statement.spec.ts

@@ -1,6 +1,6 @@
 import { expect } from 'chai';
 
-import { Dialect, identify } from '../../src';
+import { Dialect, identify, Options } from '../../src';
 
 describe('identifier', () => {
   describe('given queries with a single statement', () => {
@@ -1898,3 +1898,136 @@ describe('identifier', () => {
     });
   });
 });
+
+const options = { dialect: 'mssql', includeLeadingComments: true } as Options;
+
+describe('includeLeadingComments', () => {
+  it('keeps the default keyword offsets and text', () => {
+    const sql = '-- header\nCREATE PROCEDURE p AS SELECT 1;';
+    const actual = identify(sql, { dialect: 'mssql' })[0];
+    expect(actual.start).to.equal(sql.indexOf('CREATE'));
+    expect(actual.text).to.equal(sql.slice(actual.start, actual.end + 1));
+  });
+
+  it('starts at the first comment without leading whitespace', () => {
+    const sql = '  -- header\n/* detail */\nCREATE PROCEDURE p AS SELECT 1;';
+    const actual = identify(sql, options)[0];
+    expect(actual.start).to.equal(2);
+    expect(actual.text).to.equal(sql.slice(2));
+    expect(actual.type).to.equal('CREATE_PROCEDURE');
+  });
+
+  it('does not include empty statement delimiters', () => {
+    const sql = '; ;\n-- header\nSELECT 1;';
+    const actual = identify(sql, options)[0];
+    expect(actual.start).to.equal(4);
+    expect(actual.text).to.equal(sql.slice(4));
+  });
+
+  it('keeps each header with its own statement', () => {
+    const sql = '-- first\nCREATE VIEW a AS SELECT 1;\n/* second */\nCREATE VIEW b AS SELECT 2;';
+    const actual = identify(sql, options);
+    expect(actual).to.have.length(2);
+    expect(actual[0].text).to.equal(sql.slice(0, sql.indexOf(';') + 1));
+    expect(actual[1].start).to.equal(sql.indexOf('/* second */'));
+    expect(actual[1].text).to.equal(sql.slice(actual[1].start));
+  });
+
+  it('does not change whitespace-only prefixes', () => {
+    expect(identify('  SELECT 1;', options)[0].start).to.equal(2);
+  });
+
+  it('does not treat comment markers in strings as headers', () => {
+    const actual = identify("SELECT '-- not a header'; SELECT 2;", options);
+    expect(actual[1].text).to.equal('SELECT 2;');
+  });
+});
+
+describe('includeLeadingWhitespace', () => {
+  for (const includeLeadingComments of [false, true]) {
+    for (const includeLeadingWhitespace of [false, true]) {
+      const flags = {
+        dialect: 'mssql',
+        includeLeadingComments,
+        includeLeadingWhitespace,
+      } as Options;
+      it(`keeps independent boundaries with comments=${String(
+        includeLeadingComments,
+      )}, whitespace=${String(includeLeadingWhitespace)}`, () => {
+        const sql = '  -- header\n /* detail */\n  CREATE VIEW v AS SELECT 1;';
+        const actual = identify(sql, flags)[0];
+        const start = includeLeadingComments
+          ? includeLeadingWhitespace
+            ? 0
+            : sql.indexOf('--')
+          : includeLeadingWhitespace
+          ? sql.indexOf('\n  CREATE')
+          : sql.indexOf('CREATE');
+        expect(actual.start).to.equal(start);
+        expect(actual.text).to.equal(sql.slice(start));
+        expect(actual.end).to.equal(sql.length - 1);
+        expect(actual.type).to.equal('CREATE_VIEW');
+        expect(actual.parameters).to.deep.equal([]);
+      });
+    }
+  }
+
+  it('includes whitespace-only prefixes when enabled', () => {
+    const sql = ' \n\tSELECT 1;';
+    const actual = identify(sql, { includeLeadingWhitespace: true })[0];
+    expect(actual.start).to.equal(0);
+    expect(actual.text).to.equal(sql);
+  });
+
+  it('keeps delimiters out and preserves each contiguous prefix with both options', () => {
+    const sql = '; ;\n -- first\nSELECT 1;\n /* second */\nSELECT 2;';
+    const actual = identify(sql, { includeLeadingComments: true, includeLeadingWhitespace: true });
+    expect(actual).to.have.length(2);
+    expect(actual[0].start).to.equal(3);
+    expect(actual[0].text).to.equal(sql.slice(3, sql.indexOf('SELECT 1;') + 9));
+    expect(actual[1].start).to.equal(sql.indexOf('SELECT 1;') + 9);
+    expect(actual[1].text).to.equal(sql.slice(actual[1].start));
+  });
+});
+
+describe('leading boundary compatibility', () => {
+  it('does not cross a comment when only whitespace is enabled', () => {
+    const sql = '  /* header */SELECT 1;';
+    const actual = identify(sql, { includeLeadingWhitespace: true })[0];
+    expect(actual.start).to.equal(sql.indexOf('SELECT'));
+    expect(actual.text).to.equal('SELECT 1;');
+  });
+
+  it('does not cross a delimiter when only whitespace is enabled', () => {
+    const sql = '; ; \nSELECT 1;';
+    const actual = identify(sql, { includeLeadingWhitespace: true })[0];
+    expect(actual.start).to.equal(3);
+    expect(actual.text).to.equal(sql.slice(3));
+  });
+
+  it('keeps parameter extraction unchanged with both options enabled', () => {
+    const sql = "  -- header\nSELECT '-- literal', ?;";
+    const flags = {
+      dialect: 'mysql',
+      includeLeadingComments: true,
+      includeLeadingWhitespace: true,
+    } as Options;
+    const expected = identify(sql, { dialect: 'mysql' })[0];
+    const actual = identify(sql, flags)[0];
+    expect(actual.start).to.equal(0);
+    expect(actual.text).to.equal(sql);
+    expect(actual.end).to.equal(expected.end);
+    expect(actual.type).to.equal(expected.type);
+    expect(actual.executionType).to.equal(expected.executionType);
+    expect(actual.parameters).to.deep.equal(expected.parameters);
+  });
+});
+
+describe('line comment whitespace boundary', () => {
+  it('does not include the newline owned by an excluded line comment', () => {
+    const sql = '  -- header\n  SELECT 1;';
+    const actual = identify(sql, { includeLeadingWhitespace: true })[0];
+    expect(actual.start).to.equal(sql.indexOf('  SELECT'));
+    expect(actual.text).to.equal('  SELECT 1;');
+  });
+});