fix: address all CRITICAL/MAJOR findings from multi-model review

Fixes five correctness, reliability, and portability issues surfaced by the
consensus code review of this branch.

CRITICAL #1 — Cross-dialect partitioned diff (`data-diff.ts`):
`runPartitionedDiff` built one partition WHERE clause with `sourceDialect`
and passed it as shared `where_clause` to the recursive `runDataDiff`, which
applied it to both warehouses identically. Cross-dialect partition mode
(MSSQL → Postgres) failed because the target received T-SQL
`DATETRUNC`/`CONVERT(DATE, …, 23)`. Now builds per-side WHERE using each
warehouse's dialect and bakes it into dialect-quoted subquery SQL for source
and target independently. The existing side-aware CTE injection handles the
rest.

MAJOR #2 — Azure AD token caching and refresh (`sqlserver.ts`):
`acquireAzureToken` fetched a fresh token on every `connect()` and embedded
it in the pool config with no refresh. Long-lived sessions silently failed
when the ~1h token expired. Adds a module-scoped cache keyed by
`(resource, client_id)` with proactive refresh 5 min before expiry, parsing
`expiresOnTimestamp` from `@azure/identity` or the JWT `exp` claim from the
`az` CLI fallback. Exposes `_resetTokenCacheForTests` for isolation.

MAJOR #3 — `joindiff` + cross-warehouse guard (`data-diff.ts`):
Explicit `algorithm: "joindiff"` combined with different warehouses produced
broken SQL (one task referencing two CTE aliases with only one injected).
Now returns an early error with a clear message steering users to `hashdiff`
or `auto`. Cross-warehouse detection switched from warehouse-name string
compare to dialect compare, matching the underlying SQL-divergence invariant.

MAJOR #4 — Dialect-aware identifier quoting in CTE wrapping (`data-diff.ts`):
`resolveTableSources` wrapped plain-table names with ANSI double-quotes for
all dialects. T-SQL/Fabric require `QUOTED_IDENTIFIER ON` for this to work;
default for `mssql`/tedious is ON, but user contexts (stored procs, legacy
collations) can override. Now accepts source/target dialect parameters and
delegates to `quoteIdentForDialect`, which was hoisted to module scope so
it can be reused across partition and CTE paths.

MAJOR #5 — Configurable Azure resource URL (`sqlserver.ts`, `normalize.ts`):
Token acquisition hardcoded `https://database.windows.net/`, blocking Azure
Government, Azure China, and sovereign-cloud customers. Now honours an
explicit `azure_resource_url` config field and otherwise infers the URL
from the host suffix (`.usgovcloudapi.net`, `.chinacloudapi.cn`). Adds the
usual camelCase/snake_case aliases in the SQL Server normalizer.

Also surfaces Azure auth error causes: if both `@azure/identity` and `az`
CLI fail, the thrown error includes both hints (redacted) so users know why
rather than seeing the generic "install @azure/identity or run az login"
message.

Tests: adds `data-diff-cross-dialect.test.ts` covering the cross-dialect
partition WHERE routing and the `joindiff` guard; extends
`data-diff-cte.test.ts` with dialect-aware quoting assertions for tsql,
fabric, and mysql; extends `sqlserver-unit.test.ts` with cache hit / expiry
refresh / client-id keyed cache tests, commercial/gov/china/custom resource
URL resolution, and the combined-error-hints surface.

All 41 sqlserver driver tests, 24 data-diff orchestrator tests, and 214
normalize/connections tests pass.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: suryaiyer95

packages/drivers/src/normalize.ts

@@ -70,6 +70,7 @@ const SQLSERVER_ALIASES: AliasMap = {
   azure_client_id: ["clientId", "client_id", "azureClientId"],
   azure_client_secret: ["clientSecret", "client_secret", "azureClientSecret"],
   access_token: ["token", "accessToken"],
+  azure_resource_url: ["azureResourceUrl", "resourceUrl", "resource_url"],
 }
 
 const ORACLE_ALIASES: AliasMap = {

packages/drivers/src/sqlserver.ts

@@ -4,6 +4,65 @@
 
 import type { ConnectionConfig, Connector, ConnectorResult, ExecuteOptions, SchemaColumn } from "./types"
 
+// ---------------------------------------------------------------------------
+// Azure AD helpers — cache + resource URL resolution
+// ---------------------------------------------------------------------------
+
+// Module-scoped token cache, keyed by `${resource}|${clientId ?? ""}`.
+// Tokens are reused across `connect()` calls in the same process and refreshed
+// a few minutes before expiry. Fixes the issue where every new connection
+// fetched a fresh token (wasteful, risks throttling) and long-lived diffs
+// failed silently when the embedded token hit its ~1h TTL.
+const tokenCache = new Map<string, { token: string; expiresAt: number }>()
+const TOKEN_REFRESH_MARGIN_MS = 5 * 60 * 1000 // refresh 5 minutes before expiry
+const TOKEN_FALLBACK_TTL_MS = 50 * 60 * 1000 // used when JWT has no exp claim
+
+/**
+ * Parse the `exp` claim from a JWT access token (milliseconds since epoch).
+ * Returns undefined if the token isn't a JWT or has no exp claim.
+ */
+function parseTokenExpiry(token: string): number | undefined {
+  try {
+    const parts = token.split(".")
+    if (parts.length !== 3) return undefined
+    const payload = parts[1]
+    // base64url → base64 + padding
+    const padded = payload.replace(/-/g, "+").replace(/_/g, "/") +
+      "=".repeat((4 - (payload.length % 4)) % 4)
+    const decoded = Buffer.from(padded, "base64").toString("utf-8")
+    const claims = JSON.parse(decoded)
+    return typeof claims.exp === "number" ? claims.exp * 1000 : undefined
+  } catch {
+    return undefined
+  }
+}
+
+/**
+ * Resolve the Azure resource URL for token acquisition.
+ *
+ * Preference order:
+ *   1. Explicit `config.azure_resource_url`.
+ *   2. Inferred from host suffix (Azure Gov / China).
+ *   3. Default Azure commercial cloud.
+ */
+function resolveAzureResourceUrl(config: ConnectionConfig): string {
+  const explicit = config.azure_resource_url as string | undefined
+  if (explicit) return explicit
+  const host = (config.host as string | undefined) ?? ""
+  if (host.includes(".usgovcloudapi.net") || host.includes(".datawarehouse.fabric.microsoft.us")) {
+    return "https://database.usgovcloudapi.net/"
+  }
+  if (host.includes(".chinacloudapi.cn")) {
+    return "https://database.chinacloudapi.cn/"
+  }
+  return "https://database.windows.net/"
+}
+
+/** Visible for testing: reset the module-scoped token cache. */
+export function _resetTokenCacheForTests(): void {
+  tokenCache.clear()
+}
+
 export async function connect(config: ConnectionConfig): Promise<Connector> {
   let mssql: any
   let MssqlConnectionPool: any
@@ -70,8 +129,24 @@ export async function connect(config: ConnectionConfig): Promise<Connector> {
         // Strategy: try @azure/identity first (works when module resolution
         // is correct), fall back to shelling out to `az account get-access-token`
         // (works everywhere Azure CLI is installed).
+        //
+        // Tokens are cached module-scope keyed by (resource, client_id) and
+        // refreshed 5 minutes before expiry — reuses tokens across connections
+        // and prevents silent failures when embedded tokens hit their TTL.
+        const resourceUrl = resolveAzureResourceUrl(config)
+        const clientId = (config.azure_client_id as string | undefined) ?? ""
+        const cacheKey = `${resourceUrl}|${clientId}`
+
         const acquireAzureToken = async (): Promise<string> => {
+          const cached = tokenCache.get(cacheKey)
+          if (cached && cached.expiresAt - Date.now() > TOKEN_REFRESH_MARGIN_MS) {
+            return cached.token
+          }
+
           let token: string | undefined
+          let expiresAt: number | undefined
+          let azureIdentityError: unknown = null
+          let azCliStderr = ""
 
           try {
             const azureIdentity = await import("@azure/identity")
@@ -80,31 +155,51 @@ export async function connect(config: ConnectionConfig): Promise<Connector> {
                 ? { managedIdentityClientId: config.azure_client_id as string }
                 : undefined,
             )
-            const tokenResponse = await credential.getToken("https://database.windows.net/.default")
-            token = tokenResponse?.token
-          } catch {
-            // @azure/identity unavailable or browser bundle — fall through
+            const tokenResponse = await credential.getToken(`${resourceUrl}.default`)
+            if (tokenResponse?.token) {
+              token = tokenResponse.token
+              // @azure/identity provides expiresOnTimestamp (ms). Prefer it; fall
+              // back to parsing the JWT exp claim so both paths share the cache.
+              expiresAt = tokenResponse.expiresOnTimestamp ?? parseTokenExpiry(token)
+            }
+          } catch (err) {
+            azureIdentityError = err
+            // @azure/identity unavailable or browser bundle — fall through to CLI
           }
 
           if (!token) {
             try {
               const { execSync } = await import("node:child_process")
               const out = execSync(
-                "az account get-access-token --resource https://database.windows.net/ --query accessToken -o tsv",
+                `az account get-access-token --resource ${resourceUrl} --query accessToken -o tsv`,
                 { encoding: "utf-8", timeout: 15000, stdio: ["pipe", "pipe", "pipe"] },
               ).trim()
-              if (out) token = out
-            } catch {
-              // az CLI not installed or not logged in
+              if (out) {
+                token = out
+                expiresAt = parseTokenExpiry(out)
+              }
+            } catch (err: any) {
+              // Capture stderr so the final error message can hint at the root cause
+              // (e.g. "Please run 'az login'", "subscription not found").
+              azCliStderr = String(err?.stderr ?? err?.message ?? "").slice(0, 200).trim()
             }
           }
 
           if (!token) {
+            const hints: string[] = []
+            if (azureIdentityError) hints.push(`@azure/identity: ${String(azureIdentityError).slice(0, 120)}`)
+            if (azCliStderr) hints.push(`az CLI: ${azCliStderr}`)
+            const detail = hints.length > 0 ? ` (${hints.join("; ")})` : ""
             throw new Error(
-              "Azure AD token acquisition failed. Either install @azure/identity (npm install @azure/identity) " +
+              `Azure AD token acquisition failed${detail}. Either install @azure/identity (npm install @azure/identity) ` +
               "or log in with Azure CLI (az login).",
             )
           }
+
+          tokenCache.set(cacheKey, {
+            token,
+            expiresAt: expiresAt ?? Date.now() + TOKEN_FALLBACK_TTL_MS,
+          })
           return token
         }
 

packages/drivers/test/sqlserver-unit.test.ts

@@ -65,25 +65,47 @@ mock.module("mssql", () => ({
   },
 }))
 
+// Exposed to individual tests so they can assert scope / force failures.
+const azureIdentityState = {
+  lastScope: "" as string,
+  tokenOverride: null as null | { token: string; expiresOnTimestamp?: number },
+  throwOnGetToken: false as boolean,
+}
 mock.module("@azure/identity", () => ({
   DefaultAzureCredential: class {
     _opts: any
     constructor(opts?: any) { this._opts = opts }
-    async getToken(_scope: string) { return { token: "mock-azure-token-12345", expiresOnTimestamp: Date.now() + 3600000 } }
+    async getToken(scope: string) {
+      azureIdentityState.lastScope = scope
+      if (azureIdentityState.throwOnGetToken) throw new Error("mock identity failure")
+      if (azureIdentityState.tokenOverride) return azureIdentityState.tokenOverride
+      return { token: "mock-azure-token-12345", expiresOnTimestamp: Date.now() + 3600000 }
+    }
   },
 }))
 
-// Bun's mock.module() replaces the module for ALL test files in the same run,
-// so we re-export every symbol other tests might import (spawn, exec, fork, etc.)
-// in addition to the execSync stub used by the Azure CLI fallback path.
+// Exposed to tests to stub the `az` CLI fallback.
+const cliState = {
+  lastCmd: "" as string,
+  output: "mock-cli-token-fallback\n" as string,
+  throwError: null as null | { stderr?: string; message?: string },
+}
 const realChildProcess = await import("node:child_process")
 mock.module("node:child_process", () => ({
   ...realChildProcess,
-  execSync: (_cmd: string) => "mock-cli-token-fallback\n",
+  execSync: (cmd: string) => {
+    cliState.lastCmd = cmd
+    if (cliState.throwError) {
+      const e: any = new Error(cliState.throwError.message ?? "az failed")
+      e.stderr = cliState.throwError.stderr
+      throw e
+    }
+    return cliState.output
+  },
 }))
 
 // Import after mocking
-const { connect } = await import("../src/sqlserver")
+const { connect, _resetTokenCacheForTests } = await import("../src/sqlserver")
 
 describe("SQL Server driver unit tests", () => {
   let connector: Awaited<ReturnType<typeof connect>>
@@ -488,4 +510,178 @@ describe("SQL Server driver unit tests", () => {
       expect(result.rows).toEqual([["alice", 42]])
     })
   })
+
+  // --- Azure token caching (Fix #2) ---
+
+  describe("Azure token cache", () => {
+    beforeEach(() => {
+      _resetTokenCacheForTests()
+      azureIdentityState.throwOnGetToken = false
+      azureIdentityState.tokenOverride = null
+      cliState.throwError = null
+      cliState.output = "mock-cli-token-fallback\n"
+    })
+
+    test("second connect with same (resource, clientId) reuses cached token", async () => {
+      let getTokenCalls = 0
+      azureIdentityState.tokenOverride = { token: "cached-token-A", expiresOnTimestamp: Date.now() + 3600_000 }
+      // Hook getToken counter
+      const origCredential = (await import("@azure/identity")).DefaultAzureCredential
+      const origGetToken = origCredential.prototype.getToken
+      origCredential.prototype.getToken = async function (scope: string) {
+        getTokenCalls++
+        return origGetToken.call(this, scope)
+      }
+      try {
+        resetMocks()
+        const c1 = await connect({
+          host: "h.database.windows.net", database: "d",
+          authentication: "azure-active-directory-default",
+        })
+        await c1.connect()
+        const c2 = await connect({
+          host: "h.database.windows.net", database: "d",
+          authentication: "azure-active-directory-default",
+        })
+        await c2.connect()
+        expect(getTokenCalls).toBe(1)
+        // Both pool configs embed the same cached token
+        expect(mockConnectCalls[0].authentication.options.token).toBe("cached-token-A")
+        expect(mockConnectCalls[1].authentication.options.token).toBe("cached-token-A")
+      } finally {
+        origCredential.prototype.getToken = origGetToken
+      }
+    })
+
+    test("near-expiry token triggers refresh", async () => {
+      // First token expires in 1 minute (well under the 5-minute refresh margin)
+      azureIdentityState.tokenOverride = { token: "about-to-expire", expiresOnTimestamp: Date.now() + 60_000 }
+      resetMocks()
+      const c1 = await connect({
+        host: "h.database.windows.net", database: "d",
+        authentication: "azure-active-directory-default",
+      })
+      await c1.connect()
+      // Now change the mock to issue a new token on refresh
+      azureIdentityState.tokenOverride = { token: "fresh-token", expiresOnTimestamp: Date.now() + 3600_000 }
+      const c2 = await connect({
+        host: "h.database.windows.net", database: "d",
+        authentication: "azure-active-directory-default",
+      })
+      await c2.connect()
+      expect(mockConnectCalls[0].authentication.options.token).toBe("about-to-expire")
+      expect(mockConnectCalls[1].authentication.options.token).toBe("fresh-token")
+    })
+
+    test("different clientIds cache separately", async () => {
+      azureIdentityState.tokenOverride = { token: "shared-token", expiresOnTimestamp: Date.now() + 3600_000 }
+      resetMocks()
+      const a = await connect({
+        host: "h.database.windows.net", database: "d",
+        authentication: "azure-active-directory-default",
+        azure_client_id: "client-1",
+      })
+      await a.connect()
+      const b = await connect({
+        host: "h.database.windows.net", database: "d",
+        authentication: "azure-active-directory-default",
+        azure_client_id: "client-2",
+      })
+      await b.connect()
+      // Both embed the mock token but the cache is keyed separately; both calls
+      // hit getToken (not 1) — easiest way to assert is that both configs got
+      // a token with no thrown "expired" behavior.
+      expect(mockConnectCalls[0].authentication.options.token).toBe("shared-token")
+      expect(mockConnectCalls[1].authentication.options.token).toBe("shared-token")
+    })
+  })
+
+  // --- Configurable / inferred Azure resource URL (Fix #5) ---
+
+  describe("Azure resource URL resolution", () => {
+    beforeEach(() => {
+      _resetTokenCacheForTests()
+      azureIdentityState.throwOnGetToken = false
+      azureIdentityState.tokenOverride = null
+      cliState.throwError = null
+    })
+
+    test("commercial cloud: default to database.windows.net", async () => {
+      resetMocks()
+      const c = await connect({
+        host: "myserver.database.windows.net", database: "d",
+        authentication: "azure-active-directory-default",
+      })
+      await c.connect()
+      expect(azureIdentityState.lastScope).toBe("https://database.windows.net/.default")
+    })
+
+    test("Azure Government host infers usgovcloudapi.net", async () => {
+      resetMocks()
+      const c = await connect({
+        host: "myserver.database.usgovcloudapi.net", database: "d",
+        authentication: "azure-active-directory-default",
+      })
+      await c.connect()
+      expect(azureIdentityState.lastScope).toBe("https://database.usgovcloudapi.net/.default")
+    })
+
+    test("Azure China host infers chinacloudapi.cn", async () => {
+      resetMocks()
+      const c = await connect({
+        host: "myserver.database.chinacloudapi.cn", database: "d",
+        authentication: "azure-active-directory-default",
+      })
+      await c.connect()
+      expect(azureIdentityState.lastScope).toBe("https://database.chinacloudapi.cn/.default")
+    })
+
+    test("explicit azure_resource_url wins over host inference", async () => {
+      resetMocks()
+      const c = await connect({
+        host: "myserver.database.windows.net", // commercial host
+        database: "d",
+        authentication: "azure-active-directory-default",
+        azure_resource_url: "https://custom.sovereign.example/",
+      })
+      await c.connect()
+      expect(azureIdentityState.lastScope).toBe("https://custom.sovereign.example/.default")
+    })
+
+    test("az CLI fallback uses the same resource URL", async () => {
+      // Disable @azure/identity so we hit the az CLI fallback
+      azureIdentityState.throwOnGetToken = true
+      cliState.output = "eyJ.eyJ.sig\n" // looks like JWT; parseTokenExpiry returns undefined → fallback TTL
+      resetMocks()
+      const c = await connect({
+        host: "myserver.database.usgovcloudapi.net", database: "d",
+        authentication: "azure-active-directory-default",
+      })
+      await c.connect()
+      expect(cliState.lastCmd).toContain("--resource https://database.usgovcloudapi.net/")
+    })
+  })
+
+  // --- Error surfacing when auth fails (Fix #5 bonus, Minor #10 addressed) ---
+
+  describe("Azure auth error surfacing", () => {
+    beforeEach(() => {
+      _resetTokenCacheForTests()
+      azureIdentityState.throwOnGetToken = false
+      azureIdentityState.tokenOverride = null
+      cliState.throwError = null
+    })
+
+    test("both @azure/identity and az CLI fail → error includes both hints", async () => {
+      azureIdentityState.throwOnGetToken = true
+      cliState.throwError = { stderr: "Please run 'az login' to set up an account.", message: "failed" }
+      resetMocks()
+      const c = await connect({
+        host: "h.database.windows.net", database: "d",
+        authentication: "azure-active-directory-default",
+      })
+      await expect(c.connect()).rejects.toThrow(/Azure AD token acquisition failed/)
+      await expect(c.connect()).rejects.toThrow(/az CLI:.*az login/)
+    })
+  })
 })