From cb81b2ef2b4b909fbac2222cc727f02bfaaf743d Mon Sep 17 00:00:00 2001
From: Mikhail Cheshkov <mcheshkov@gmail.com>
Date: Wed, 26 Feb 2025 01:34:30 +0200
Subject: [PATCH] test: Add multidb SQL pushdown smoke test

---
 .../__snapshots__/smoke-multidb.test.ts.snap  | 16 +++++
 .../cubejs-testing/test/smoke-multidb.test.ts | 64 +++++++++++++++++++
 2 files changed, 80 insertions(+)

diff --git a/packages/cubejs-testing/test/__snapshots__/smoke-multidb.test.ts.snap b/packages/cubejs-testing/test/__snapshots__/smoke-multidb.test.ts.snap
index b1bb0e431e33e..1d8d70145453e 100644
--- a/packages/cubejs-testing/test/__snapshots__/smoke-multidb.test.ts.snap
+++ b/packages/cubejs-testing/test/__snapshots__/smoke-multidb.test.ts.snap
@@ -1,5 +1,21 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
+exports[`multidb SQL pushdown queries to different data sources: Products 1`] = `
+Array [
+  Object {
+    "name": "apples",
+  },
+]
+`;
+
+exports[`multidb SQL pushdown queries to different data sources: Suppliers 1`] = `
+Array [
+  Object {
+    "company": "Fruits Inc",
+  },
+]
+`;
+
 exports[`multidb query: query 1`] = `
 Array [
   Object {
diff --git a/packages/cubejs-testing/test/smoke-multidb.test.ts b/packages/cubejs-testing/test/smoke-multidb.test.ts
index b5f5df5295c86..4e869a20cf2f5 100644
--- a/packages/cubejs-testing/test/smoke-multidb.test.ts
+++ b/packages/cubejs-testing/test/smoke-multidb.test.ts
@@ -3,6 +3,7 @@ import { MysqlDBRunner, PostgresDBRunner } from '@cubejs-backend/testing-shared'
 import cubejs, { CubeApi } from '@cubejs-client/core';
 // eslint-disable-next-line import/no-extraneous-dependencies
 import { afterAll, beforeAll, expect, jest } from '@jest/globals';
+import { Client as PgClient } from 'pg';
 import { BirdBox, getBirdbox } from '../src';
 import {
   DEFAULT_API_TOKEN,
@@ -11,12 +12,43 @@ import {
   JEST_BEFORE_ALL_DEFAULT_TIMEOUT,
 } from './smoke-tests';
 
+// TODO: Random port?
+const pgPort = 5656;
+let connectionId = 0;
+
+async function createPostgresClient(user: string, password: string) {
+  connectionId++;
+  const currentConnId = connectionId;
+
+  console.debug(`[pg] new connection ${currentConnId}`);
+
+  const conn = new PgClient({
+    database: 'db',
+    port: pgPort,
+    host: '127.0.0.1',
+    user,
+    password,
+    ssl: false,
+  });
+  conn.on('error', (err) => {
+    console.log(err);
+  });
+  conn.on('end', () => {
+    console.debug(`[pg] end ${currentConnId}`);
+  });
+
+  await conn.connect();
+
+  return conn;
+}
+
 describe('multidb', () => {
   jest.setTimeout(60 * 5 * 1000);
   let db: StartedTestContainer;
   let db2: StartedTestContainer;
   let birdbox: BirdBox;
   let client: CubeApi;
+  let connection: PgClient;
 
   beforeAll(async () => {
     db = await PostgresDBRunner.startContainer({});
@@ -39,6 +71,9 @@ describe('multidb', () => {
         CUBEJS_DB_USER2: 'root',
         CUBEJS_DB_PASS2: 'Test1test',
 
+        CUBEJS_PG_SQL_PORT: `${pgPort}`,
+        CUBESQL_SQL_PUSH_DOWN: 'true',
+
         ...DEFAULT_CONFIG,
       },
       {
@@ -49,6 +84,7 @@ describe('multidb', () => {
     client = cubejs(async () => DEFAULT_API_TOKEN, {
       apiUrl: birdbox.configuration.apiUrl,
     });
+    connection = await createPostgresClient('admin', 'admin_password');
   }, JEST_BEFORE_ALL_DEFAULT_TIMEOUT);
 
   afterAll(async () => {
@@ -69,4 +105,32 @@ describe('multidb', () => {
     });
     expect(response.rawData()).toMatchSnapshot('query');
   });
+
+  test('SQL pushdown queries to different data sources: Products', async () => {
+    const res = await connection.query(`
+  SELECT
+    name
+  FROM
+    Products
+  WHERE
+    LOWER(name) = 'apples'
+  GROUP BY
+    1
+    `);
+    expect(res.rows).toMatchSnapshot();
+  });
+
+  test('SQL pushdown queries to different data sources: Suppliers', async () => {
+    const res = await connection.query(`
+  SELECT
+    company
+  FROM
+    Suppliers
+  WHERE
+    LOWER(company) = 'fruits inc'
+  GROUP BY
+    1
+    `);
+    expect(res.rows).toMatchSnapshot();
+  });
 });