@lightdash/warehouses 0.3414.0 → 0.3416.0

package/dist/warehouseClients/BigqueryWarehouseClient.d.ts

@@ -1,6 +1,6 @@
-import { BigQuery, Dataset, QueryRowsResponse } from '@google-cloud/bigquery';
+import { BigQuery, Dataset, DatasetsResponse, GetDatasetsOptions, QueryRowsResponse } from '@google-cloud/bigquery';
 import bigquery from '@google-cloud/bigquery/build/src/types';
-import { AnyType, BigqueryDataset, BigqueryProject, CreateBigqueryCredentials, DimensionType, Metric, PartitionColumn, SupportedDbtAdapter, TimeIntervalUnit, WarehouseQueryError, WarehouseResults, WarehouseTypes } from '@lightdash/common';
+import { AnyType, BigqueryDataset, BigqueryProject, BigqueryProjectRecommendation, CreateBigqueryCredentials, DimensionType, Metric, PartitionColumn, SupportedDbtAdapter, TimeIntervalUnit, WarehouseQueryError, WarehouseResults, WarehouseTypes } from '@lightdash/common';
 import { WarehouseCatalog, WarehouseExecuteAsyncQuery, WarehouseExecuteAsyncQueryArgs } from '../types';
 import WarehouseBaseClient from './WarehouseBaseClient';
 import WarehouseBaseSqlBuilder from './WarehouseBaseSqlBuilder';
@@ -90,7 +90,17 @@ export declare class BigqueryWarehouseClient extends WarehouseBaseClient<CreateB
     parseError(error: bigquery.IErrorProto, query?: string): WarehouseQueryError;
     executeAsyncQuery({ sql, tags, timezone }: WarehouseExecuteAsyncQueryArgs, resultsStreamCallback?: (rows: WarehouseResults['rows'], fields: WarehouseResults['fields']) => void): Promise<WarehouseExecuteAsyncQuery>;
     private awaitJobCompletion;
-    static getDatabases(projectId: string, refresh_token: string): Promise<BigqueryDataset[]>;
+    static getDatabases(projectId: string, refresh_token: string, bigqueryClient?: {
+        getDatasets: () => Promise<DatasetsResponse>;
+        query: (query: string) => Promise<QueryRowsResponse>;
+    }): Promise<BigqueryDataset[]>;
+    static getProjectRecommendation(projects: BigqueryProject[], refreshToken: string, bigqueryClientFactory?: (projectId: string) => {
+        getDatasets: (options: GetDatasetsOptions) => Promise<DatasetsResponse>;
+        query: (options: {
+            query: string;
+            location: string;
+        }) => Promise<QueryRowsResponse>;
+    }): Promise<BigqueryProjectRecommendation>;
     static getProjects(accessToken: string): Promise<BigqueryProject[]>;
 }
 export {};

package/dist/warehouseClients/BigqueryWarehouseClient.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"BigqueryWarehouseClient.d.ts","sourceRoot":"","sources":["../../src/warehouseClients/BigqueryWarehouseClient.ts"],"names":[],"mappings":"AAAA,OAAO,EACH,QAAQ,EAKR,OAAO,EAGP,iBAAiB,EACpB,MAAM,wBAAwB,CAAC;AAChC,OAAO,QAAQ,MAAM,wCAAwC,CAAC;AAC9D,OAAO,EACH,OAAO,EAEP,eAAe,EACf,eAAe,EACf,yBAAyB,EACzB,aAAa,EAEb,MAAM,EAEN,eAAe,EAIf,mBAAmB,EACnB,gBAAgB,EAEhB,mBAAmB,EACnB,gBAAgB,EAChB,cAAc,EACjB,MAAM,mBAAmB,CAAC;AAG3B,OAAO,EACH,gBAAgB,EAChB,0BAA0B,EAC1B,8BAA8B,EAEjC,MAAM,UAAU,CAAC;AAMlB,OAAO,mBAAmB,MAAM,uBAAuB,CAAC;AACxD,OAAO,uBAAuB,MAAM,2BAA2B,CAAC;AAEhE,oBAAY,iBAAiB;IACzB,MAAM,WAAW;IACjB,OAAO,YAAY;IACnB,KAAK,UAAU;IACf,KAAK,UAAU;IACf,KAAK,UAAU;IACf,OAAO,YAAY;IACnB,OAAO,YAAY;IACnB,IAAI,SAAS;IACb,SAAS,cAAc;IACvB,IAAI,SAAS;IACb,IAAI,SAAS;IACb,QAAQ,aAAa;IACrB,SAAS,cAAc;IACvB,OAAO,YAAY;IACnB,UAAU,eAAe;IACzB,MAAM,WAAW;IACjB,MAAM,WAAW;IACjB,KAAK,UAAU;CAClB;AAsDD,KAAK,WAAW,GAAG;IACf,MAAM,EAAE,YAAY,EAAE,CAAC;CAC1B,CAAC;AAEF,KAAK,YAAY,GAAG,QAAQ,CAAC,IAAI,CAAC,QAAQ,CAAC,iBAAiB,EAAE,MAAM,GAAG,MAAM,CAAC,CAAC,CAAC;AAehF,KAAK,aAAa,GAAG;IACjB,MAAM,EAAE,QAAQ,CAAC,WAAW,EAAE,CAAC;CAClC,CAAC;AAEF,qBAAa,kBAAmB,SAAQ,uBAAuB;IAC3D,QAAQ,CAAC,IAAI,2BAA2B;IAExC,cAAc,IAAI,mBAAmB;IAIrC,YAAY,CAAC,GAAG,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,MAAM;IAajD,iBAAiB,IAAI,MAAM;IAI3B,eAAe,IAAI,MAAM;IAIzB,YAAY,CAAC,KAAK,EAAE,MAAM,GAAG,MAAM;IAmBnC,eAAe,CAAC,IAAI,EAAE,IAAI,GAAG,MAAM;IAKnC,cAAc,CAAC,KAAK,EAAE,MAAM,EAAE,IAAI,EAAE,gBAAgB,GAAG,MAAM;IAM7D,uBAAuB,CACnB,iBAAiB,EAAE,MAAM,EACzB,eAAe,EAAE,MAAM,GACxB,MAAM;IAKT,YAAY,CAAC,QAAQ,EAAE,MAAM,GAAG,MAAM;IAKtC,UAAU,CAAC,QAAQ,EAAE,MAAM,EAAE,GAAG,MAAM;IAKtC,aAAa,CAAC,UAAU,EAAE,MAAM,EAAE,OAAO,CAAC,EAAE,MAAM,GAAG,MAAM;CAO9D;AAED,qBAAa,uBAAwB,SAAQ,mBAAmB,CAAC,yBAAyB,CAAC;IACvF,OAAO,CAAC,MAAM,CAAC,QAAQ,CAAC,UAAU,CAAM;IAExC,MAAM,EAAE,QAAQ,CAAC;gBAEL,WAAW,EAAE,yBAAyB;IA2BlD,MAAM,CAAC,eAAe,CAAC,KAAK,EAAE,OAAO,GAAG,KAAK,IAAI,aAAa;IAI9D;;;;;;;;;;OAUG;IACH,MAAM,CAAC,wBAAwB,CAC3B,MAAM,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,GAChC,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,GAAG,SAAS;IAgCrC,MAAM,CAAC,qBAAqB,CAAC,QAAQ,EAAE,iBAAiB,CAAC,CAAC,CAAC,GAAG,SAAS;cAExC,aAAa;;YAY9B,SAAS;YA+BT,MAAM;YAWN,qBAAqB;YAUrB,aAAa;IAoCrB,WAAW,CACb,KAAK,EAAE,MAAM,EACb,cAAc,EAAE,CAAC,IAAI,EAAE,gBAAgB,KAAK,IAAI,GAAG,OAAO,CAAC,IAAI,CAAC,EAChE,OAAO,EAAE;QACL,MAAM,CAAC,EAAE,OAAO,EAAE,CAAC;QACnB,IAAI,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;QAC9B,QAAQ,CAAC,EAAE,MAAM,CAAC;KACrB,GACF,OAAO,CAAC,IAAI,CAAC;WAwBH,gBAAgB,CACzB,OAAO,EAAE,OAAO,EAChB,KAAK,EAAE,MAAM,GACd,OAAO,CAAC,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,WAAW,CAAC,CAAC;IAW3C,UAAU,CACZ,QAAQ,EAAE;QACN,QAAQ,EAAE,MAAM,CAAC;QACjB,MAAM,EAAE,MAAM,CAAC;QACf,KAAK,EAAE,MAAM,CAAC;KACjB,EAAE;IA6CD,YAAY;;;;;;IA+DZ,SAAS,CACX,SAAS,EAAE,MAAM,EACjB,MAAM,EAAE,MAAM,EACd,QAAQ,CAAC,EAAE,MAAM,GAClB,OAAO,CAAC,gBAAgB,CAAC;IAoB5B,UAAU,CAAC,KAAK,EAAE,QAAQ,CAAC,WAAW,EAAE,KAAK,GAAE,MAAW;IAiGpD,iBAAiB,CACnB,EAAE,GAAG,EAAE,IAAI,EAAE,QAAQ,EAAE,EAAE,8BAA8B,EACvD,qBAAqB,CAAC,EAAE,CACpB,IAAI,EAAE,gBAAgB,CAAC,MAAM,CAAC,EAC9B,MAAM,EAAE,gBAAgB,CAAC,QAAQ,CAAC,KACjC,IAAI,GACV,OAAO,CAAC,0BAA0B,CAAC;YAgExB,kBAAkB;WAWnB,YAAY,CACrB,SAAS,EAAE,MAAM,EACjB,aAAa,EAAE,MAAM,GACtB,OAAO,CAAC,eAAe,EAAE,CAAC;WAoBhB,WAAW,CAAC,WAAW,EAAE,MAAM,GAAG,OAAO,CAAC,eAAe,EAAE,CAAC;CA2D5E"}
+{"version":3,"file":"BigqueryWarehouseClient.d.ts","sourceRoot":"","sources":["../../src/warehouseClients/BigqueryWarehouseClient.ts"],"names":[],"mappings":"AAAA,OAAO,EACH,QAAQ,EAKR,OAAO,EACP,gBAAgB,EAChB,kBAAkB,EAGlB,iBAAiB,EACpB,MAAM,wBAAwB,CAAC;AAChC,OAAO,QAAQ,MAAM,wCAAwC,CAAC;AAC9D,OAAO,EACH,OAAO,EAEP,eAAe,EACf,eAAe,EACf,6BAA6B,EAC7B,yBAAyB,EACzB,aAAa,EAEb,MAAM,EAEN,eAAe,EAIf,mBAAmB,EACnB,gBAAgB,EAEhB,mBAAmB,EACnB,gBAAgB,EAChB,cAAc,EACjB,MAAM,mBAAmB,CAAC;AAG3B,OAAO,EACH,gBAAgB,EAChB,0BAA0B,EAC1B,8BAA8B,EAEjC,MAAM,UAAU,CAAC;AAMlB,OAAO,mBAAmB,MAAM,uBAAuB,CAAC;AACxD,OAAO,uBAAuB,MAAM,2BAA2B,CAAC;AAKhE,oBAAY,iBAAiB;IACzB,MAAM,WAAW;IACjB,OAAO,YAAY;IACnB,KAAK,UAAU;IACf,KAAK,UAAU;IACf,KAAK,UAAU;IACf,OAAO,YAAY;IACnB,OAAO,YAAY;IACnB,IAAI,SAAS;IACb,SAAS,cAAc;IACvB,IAAI,SAAS;IACb,IAAI,SAAS;IACb,QAAQ,aAAa;IACrB,SAAS,cAAc;IACvB,OAAO,YAAY;IACnB,UAAU,eAAe;IACzB,MAAM,WAAW;IACjB,MAAM,WAAW;IACjB,KAAK,UAAU;CAClB;AAsDD,KAAK,WAAW,GAAG;IACf,MAAM,EAAE,YAAY,EAAE,CAAC;CAC1B,CAAC;AAEF,KAAK,YAAY,GAAG,QAAQ,CAAC,IAAI,CAAC,QAAQ,CAAC,iBAAiB,EAAE,MAAM,GAAG,MAAM,CAAC,CAAC,CAAC;AAehF,KAAK,aAAa,GAAG;IACjB,MAAM,EAAE,QAAQ,CAAC,WAAW,EAAE,CAAC;CAClC,CAAC;AAEF,qBAAa,kBAAmB,SAAQ,uBAAuB;IAC3D,QAAQ,CAAC,IAAI,2BAA2B;IAExC,cAAc,IAAI,mBAAmB;IAIrC,YAAY,CAAC,GAAG,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,MAAM;IAajD,iBAAiB,IAAI,MAAM;IAI3B,eAAe,IAAI,MAAM;IAIzB,YAAY,CAAC,KAAK,EAAE,MAAM,GAAG,MAAM;IAmBnC,eAAe,CAAC,IAAI,EAAE,IAAI,GAAG,MAAM;IAKnC,cAAc,CAAC,KAAK,EAAE,MAAM,EAAE,IAAI,EAAE,gBAAgB,GAAG,MAAM;IAM7D,uBAAuB,CACnB,iBAAiB,EAAE,MAAM,EACzB,eAAe,EAAE,MAAM,GACxB,MAAM;IAKT,YAAY,CAAC,QAAQ,EAAE,MAAM,GAAG,MAAM;IAKtC,UAAU,CAAC,QAAQ,EAAE,MAAM,EAAE,GAAG,MAAM;IAKtC,aAAa,CAAC,UAAU,EAAE,MAAM,EAAE,OAAO,CAAC,EAAE,MAAM,GAAG,MAAM;CAO9D;AAED,qBAAa,uBAAwB,SAAQ,mBAAmB,CAAC,yBAAyB,CAAC;IACvF,OAAO,CAAC,MAAM,CAAC,QAAQ,CAAC,UAAU,CAAM;IAExC,MAAM,EAAE,QAAQ,CAAC;gBAEL,WAAW,EAAE,yBAAyB;IA2BlD,MAAM,CAAC,eAAe,CAAC,KAAK,EAAE,OAAO,GAAG,KAAK,IAAI,aAAa;IAI9D;;;;;;;;;;OAUG;IACH,MAAM,CAAC,wBAAwB,CAC3B,MAAM,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,GAChC,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,GAAG,SAAS;IAgCrC,MAAM,CAAC,qBAAqB,CAAC,QAAQ,EAAE,iBAAiB,CAAC,CAAC,CAAC,GAAG,SAAS;cAExC,aAAa;;YAY9B,SAAS;YA+BT,MAAM;YAWN,qBAAqB;YAUrB,aAAa;IAoCrB,WAAW,CACb,KAAK,EAAE,MAAM,EACb,cAAc,EAAE,CAAC,IAAI,EAAE,gBAAgB,KAAK,IAAI,GAAG,OAAO,CAAC,IAAI,CAAC,EAChE,OAAO,EAAE;QACL,MAAM,CAAC,EAAE,OAAO,EAAE,CAAC;QACnB,IAAI,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;QAC9B,QAAQ,CAAC,EAAE,MAAM,CAAC;KACrB,GACF,OAAO,CAAC,IAAI,CAAC;WAwBH,gBAAgB,CACzB,OAAO,EAAE,OAAO,EAChB,KAAK,EAAE,MAAM,GACd,OAAO,CAAC,CAAC,MAAM,EAAE,MAAM,EAAE,MAAM,EAAE,WAAW,CAAC,CAAC;IAW3C,UAAU,CACZ,QAAQ,EAAE;QACN,QAAQ,EAAE,MAAM,CAAC;QACjB,MAAM,EAAE,MAAM,CAAC;QACf,KAAK,EAAE,MAAM,CAAC;KACjB,EAAE;IA6CD,YAAY;;;;;;IA+DZ,SAAS,CACX,SAAS,EAAE,MAAM,EACjB,MAAM,EAAE,MAAM,EACd,QAAQ,CAAC,EAAE,MAAM,GAClB,OAAO,CAAC,gBAAgB,CAAC;IAoB5B,UAAU,CAAC,KAAK,EAAE,QAAQ,CAAC,WAAW,EAAE,KAAK,GAAE,MAAW;IAiGpD,iBAAiB,CACnB,EAAE,GAAG,EAAE,IAAI,EAAE,QAAQ,EAAE,EAAE,8BAA8B,EACvD,qBAAqB,CAAC,EAAE,CACpB,IAAI,EAAE,gBAAgB,CAAC,MAAM,CAAC,EAC9B,MAAM,EAAE,gBAAgB,CAAC,QAAQ,CAAC,KACjC,IAAI,GACV,OAAO,CAAC,0BAA0B,CAAC;YAgExB,kBAAkB;WAWnB,YAAY,CACrB,SAAS,EAAE,MAAM,EACjB,aAAa,EAAE,MAAM,EACrB,cAAc,GAAE;QACZ,WAAW,EAAE,MAAM,OAAO,CAAC,gBAAgB,CAAC,CAAC;QAC7C,KAAK,EAAE,CAAC,KAAK,EAAE,MAAM,KAAK,OAAO,CAAC,iBAAiB,CAAC,CAAC;KASvD,GACH,OAAO,CAAC,eAAe,EAAE,CAAC;WA4BhB,wBAAwB,CACjC,QAAQ,EAAE,eAAe,EAAE,EAC3B,YAAY,EAAE,MAAM,EACpB,qBAAqB,GAAE,CAAC,SAAS,EAAE,MAAM,KAAK;QAC1C,WAAW,EAAE,CACT,OAAO,EAAE,kBAAkB,KAC1B,OAAO,CAAC,gBAAgB,CAAC,CAAC;QAC/B,KAAK,EAAE,CAAC,OAAO,EAAE;YACb,KAAK,EAAE,MAAM,CAAC;YACd,QAAQ,EAAE,MAAM,CAAC;SACpB,KAAK,OAAO,CAAC,iBAAiB,CAAC,CAAC;KAgBpC,GACF,OAAO,CAAC,6BAA6B,CAAC;WAyF5B,WAAW,CAAC,WAAW,EAAE,MAAM,GAAG,OAAO,CAAC,eAAe,EAAE,CAAC;CA2D5E"}

package/dist/warehouseClients/BigqueryWarehouseClient.js

@@ -12,6 +12,8 @@ const processPromisesInBatches_1 = require("../utils/processPromisesInBatches");
 const sql_1 = require("../utils/sql");
 const WarehouseBaseClient_1 = __importDefault(require("./WarehouseBaseClient"));
 const WarehouseBaseSqlBuilder_1 = __importDefault(require("./WarehouseBaseSqlBuilder"));
+const MAX_RECOMMENDATION_PROJECTS = 8;
+const MAX_RECOMMENDATION_REGIONS_PER_PROJECT = 3;
 var BigqueryFieldType;
 (function (BigqueryFieldType) {
     BigqueryFieldType["STRING"] = "STRING";
@@ -516,23 +518,112 @@ class BigqueryWarehouseClient extends WarehouseBaseClient_1.default {
             });
         });
     }
-    static async getDatabases(projectId, refresh_token) {
-        const bigqueryClient = new bigquery_1.BigQuery({
+    static async getDatabases(projectId, refresh_token, bigqueryClient = new bigquery_1.BigQuery({
+        projectId,
+        credentials: {
+            type: 'authorized_user',
+            client_id: process.env.AUTH_GOOGLE_OAUTH2_CLIENT_ID,
+            client_secret: process.env.AUTH_GOOGLE_OAUTH2_CLIENT_SECRET,
+            refresh_token,
+        },
+    })) {
+        const datasets = await bigqueryClient.getDatasets();
+        const databases = await Promise.all(datasets[0].map(async (dataset, index) => {
+            let sizeBytes = null;
+            if (index < 25) {
+                try {
+                    const [rows] = await bigqueryClient.query(`SELECT SUM(size_bytes) AS sizeBytes FROM \`${dataset.projectId}.${dataset.id}.__TABLES__\``);
+                    const size = Number(rows[0]?.sizeBytes);
+                    sizeBytes = Number.isFinite(size) ? size : null;
+                }
+                catch {
+                    sizeBytes = null;
+                }
+            }
+            return {
+                projectId: dataset.projectId,
+                location: dataset.location,
+                datasetId: dataset.id,
+                sizeBytes,
+            };
+        }));
+        return databases;
+    }
+    static async getProjectRecommendation(projects, refreshToken, bigqueryClientFactory = (projectId) => {
+        const client = new bigquery_1.BigQuery({
             projectId,
             credentials: {
                 type: 'authorized_user',
                 client_id: process.env.AUTH_GOOGLE_OAUTH2_CLIENT_ID,
                 client_secret: process.env.AUTH_GOOGLE_OAUTH2_CLIENT_SECRET,
-                refresh_token,
+                refresh_token: refreshToken,
             },
         });
-        const datasets = await bigqueryClient.getDatasets();
-        const databases = datasets[0].map((d) => ({
-            projectId: d.projectId,
-            location: d.location,
-            datasetId: d.id,
+        return {
+            getDatasets: (options) => client.getDatasets(options),
+            query: (options) => client.query(options.query, { location: options.location }),
+        };
+    }) {
+        const projectScores = await Promise.all(projects
+            .slice(0, MAX_RECOMMENDATION_PROJECTS)
+            .map(async ({ projectId }) => {
+            try {
+                const client = bigqueryClientFactory(projectId);
+                const [datasets] = await client.getDatasets({
+                    autoPaginate: false,
+                    maxResults: 1000,
+                });
+                const locationsByQualifier = new Map();
+                datasets.forEach((dataset) => {
+                    if (dataset.location) {
+                        locationsByQualifier.set(dataset.location.toLowerCase(), dataset.location);
+                    }
+                });
+                const locations = Array.from(locationsByQualifier.entries()).slice(0, MAX_RECOMMENDATION_REGIONS_PER_PROJECT);
+                const regionScores = await Promise.all(locations.map(async ([qualifier, location]) => {
+                    try {
+                        const [rows] = await client.query({
+                            query: `SELECT table_schema, SUM(total_logical_bytes) AS total_bytes FROM \`${projectId}\`.\`region-${qualifier}\`.INFORMATION_SCHEMA.TABLE_STORAGE GROUP BY table_schema`,
+                            location,
+                        });
+                        let largestDatasetSize = null;
+                        rows.forEach((row) => {
+                            const size = Number(row.total_bytes);
+                            if (Number.isFinite(size) &&
+                                (largestDatasetSize === null ||
+                                    size > largestDatasetSize)) {
+                                largestDatasetSize = size;
+                            }
+                        });
+                        return {
+                            error: false,
+                            sizeBytes: largestDatasetSize,
+                        };
+                    }
+                    catch {
+                        return { error: true, sizeBytes: null };
+                    }
+                }));
+                if (regionScores.some(({ error }) => error)) {
+                    return { projectId, sizeBytes: null };
+                }
+                const sizeBytes = regionScores.reduce((largestSize, regionScore) => regionScore.sizeBytes !== null &&
+                    (largestSize === null ||
+                        regionScore.sizeBytes > largestSize)
+                    ? regionScore.sizeBytes
+                    : largestSize, null);
+                return { projectId, sizeBytes };
+            }
+            catch {
+                return { projectId, sizeBytes: null };
+            }
         }));
-        return databases;
+        const recommendation = projectScores.reduce((largestProject, projectScore) => projectScore.sizeBytes !== null &&
+            (largestProject.sizeBytes === null ||
+                projectScore.sizeBytes > largestProject.sizeBytes)
+            ? projectScore
+            : largestProject, { projectId: null, sizeBytes: null });
+        return { projectId: recommendation.projectId };
     }
     static async getProjects(accessToken) {
         // Use BigQuery REST API directly instead of the SDK because:

package/dist/warehouseClients/BigqueryWarehouseClient.test.js

@@ -42,6 +42,172 @@ describe('BigqueryWarehouseClient', () => {
         expect(BigqueryWarehouseClient_mock_1.getTableResponse.getMetadata).toHaveBeenCalledTimes(1);
     });
 });
+describe('BigqueryWarehouseClient.getDatabases', () => {
+    const createDataset = (datasetId) => new bigquery_1.Dataset(new bigquery_1.BigQuery({ projectId: 'test-project' }), datasetId, {
+        location: 'EU',
+    });
+    const createClient = (datasets, query) => ({
+        getDatasets: vi
+            .fn()
+            .mockResolvedValue([datasets]),
+        query: vi.fn(query),
+    });
+    it('includes the stored size for each dataset', async () => {
+        const datasets = [createDataset('small'), createDataset('large')];
+        const queryMock = vi
+            .fn()
+            .mockResolvedValueOnce([[{ sizeBytes: 100 }]])
+            .mockResolvedValueOnce([[{ sizeBytes: 500 }]]);
+        const client = createClient(datasets, queryMock);
+        await expect(BigqueryWarehouseClient_1.BigqueryWarehouseClient.getDatabases('test-project', 'refresh-token', client)).resolves.toEqual([
+            {
+                projectId: 'test-project',
+                datasetId: 'small',
+                location: 'EU',
+                sizeBytes: 100,
+            },
+            {
+                projectId: 'test-project',
+                datasetId: 'large',
+                location: 'EU',
+                sizeBytes: 500,
+            },
+        ]);
+        expect(queryMock).toHaveBeenNthCalledWith(1, 'SELECT SUM(size_bytes) AS sizeBytes FROM `test-project.small.__TABLES__`');
+    });
+    it('returns null when an individual dataset size query fails', async () => {
+        const datasets = [createDataset('unavailable'), createDataset('ok')];
+        const queryMock = vi
+            .fn()
+            .mockRejectedValueOnce(new Error('permission denied'))
+            .mockResolvedValueOnce([[{ sizeBytes: 500 }]]);
+        const client = createClient(datasets, queryMock);
+        const result = await BigqueryWarehouseClient_1.BigqueryWarehouseClient.getDatabases('test-project', 'refresh-token', client);
+        expect(result.map(({ datasetId, sizeBytes }) => ({
+            datasetId,
+            sizeBytes,
+        }))).toEqual([
+            { datasetId: 'unavailable', sizeBytes: null },
+            { datasetId: 'ok', sizeBytes: 500 },
+        ]);
+    });
+    it('only queries sizes for the first 25 datasets', async () => {
+        const datasets = Array.from({ length: 30 }, (_, index) => createDataset(`dataset_${index}`));
+        const queryMock = vi
+            .fn()
+            .mockResolvedValue([[{ sizeBytes: 100 }]]);
+        const client = createClient(datasets, queryMock);
+        const result = await BigqueryWarehouseClient_1.BigqueryWarehouseClient.getDatabases('test-project', 'refresh-token', client);
+        expect(queryMock).toHaveBeenCalledTimes(25);
+        expect(result.slice(0, 25).every(({ sizeBytes }) => sizeBytes === 100)).toBe(true);
+        expect(result.slice(25).every(({ sizeBytes }) => sizeBytes === null)).toBe(true);
+    });
+});
+describe('BigqueryWarehouseClient.getProjectRecommendation', () => {
+    const createDataset = (projectId, datasetId, location) => new bigquery_1.Dataset(new bigquery_1.BigQuery({ projectId }), datasetId, {
+        location,
+    });
+    const createClient = (projectId, locations, queryImplementation) => {
+        const datasets = locations.map((location, index) => createDataset(projectId, `dataset_${index}`, location));
+        const getDatasets = vi.fn();
+        vi.mocked(getDatasets).mockResolvedValue([datasets]);
+        const query = vi.fn();
+        vi.mocked(query).mockImplementation(queryImplementation);
+        return { getDatasets, query };
+    };
+    it('ranks projects by their largest dataset across distinct regions', async () => {
+        const firstClient = createClient('first-project', ['EU', 'US', 'EU'], async ({ location }) => location === 'US'
+            ? [[{ table_schema: 'largest', total_bytes: 1200 }]]
+            : [[{ table_schema: 'smaller', total_bytes: 500 }]]);
+        const secondClient = createClient('second-project', ['EU'], async () => [
+            [
+                { table_schema: 'small', total_bytes: 100 },
+                { table_schema: 'large', total_bytes: 900 },
+            ],
+        ]);
+        const clients = new Map([
+            ['first-project', firstClient],
+            ['second-project', secondClient],
+        ]);
+        const clientFactory = vi.fn((projectId) => {
+            const client = clients.get(projectId);
+            if (!client)
+                throw new Error('Unexpected project');
+            return client;
+        });
+        const projects = [
+            { projectId: 'first-project', friendlyName: null },
+            { projectId: 'second-project', friendlyName: null },
+        ];
+        await expect(BigqueryWarehouseClient_1.BigqueryWarehouseClient.getProjectRecommendation(projects, 'refresh-token', clientFactory)).resolves.toEqual({ projectId: 'first-project' });
+        expect(vi.mocked(firstClient.getDatasets)).toHaveBeenCalledWith({
+            autoPaginate: false,
+            maxResults: 1000,
+        });
+        expect(vi.mocked(firstClient.query)).toHaveBeenCalledTimes(2);
+        expect(vi.mocked(firstClient.query)).toHaveBeenCalledWith({
+            query: 'SELECT table_schema, SUM(total_logical_bytes) AS total_bytes FROM `first-project`.`region-us`.INFORMATION_SCHEMA.TABLE_STORAGE GROUP BY table_schema',
+            location: 'US',
+        });
+    });
+    it('skips a project when any of its region queries errors', async () => {
+        const failingClient = createClient('failing-project', ['EU', 'US'], async ({ location }) => {
+            if (location === 'US')
+                throw new Error('permission denied');
+            return [[{ table_schema: 'largest', total_bytes: 5000 }]];
+        });
+        const healthyClient = createClient('healthy-project', ['EU'], async () => [[{ table_schema: 'available', total_bytes: 100 }]]);
+        const clients = new Map([
+            ['failing-project', failingClient],
+            ['healthy-project', healthyClient],
+        ]);
+        const clientFactory = vi.fn((projectId) => {
+            const client = clients.get(projectId);
+            if (!client)
+                throw new Error('Unexpected project');
+            return client;
+        });
+        await expect(BigqueryWarehouseClient_1.BigqueryWarehouseClient.getProjectRecommendation([
+            { projectId: 'failing-project', friendlyName: null },
+            { projectId: 'healthy-project', friendlyName: null },
+        ], 'refresh-token', clientFactory)).resolves.toEqual({ projectId: 'healthy-project' });
+    });
+    it('returns no recommendation when every project score is unavailable', async () => {
+        const client = createClient('unavailable-project', ['EU'], async () => {
+            throw new Error('permission denied');
+        });
+        const clientFactory = vi.fn();
+        vi.mocked(clientFactory).mockReturnValue(client);
+        await expect(BigqueryWarehouseClient_1.BigqueryWarehouseClient.getProjectRecommendation([{ projectId: 'unavailable-project', friendlyName: null }], 'refresh-token', clientFactory)).resolves.toEqual({ projectId: null });
+    });
+    it('caps ranking at eight projects and three regions per project', async () => {
+        const projects = Array.from({ length: 10 }, (_, index) => ({
+            projectId: `project-${index}`,
+            friendlyName: null,
+        }));
+        const clients = new Map(projects.map(({ projectId }, index) => [
+            projectId,
+            createClient(projectId, ['EU', 'US', 'asia-east1', 'us-west1'], async () => [
+                [{ table_schema: 'dataset', total_bytes: index }],
+            ]),
+        ]));
+        const clientFactory = vi.fn((projectId) => {
+            const client = clients.get(projectId);
+            if (!client)
+                throw new Error('Unexpected project');
+            return client;
+        });
+        await expect(BigqueryWarehouseClient_1.BigqueryWarehouseClient.getProjectRecommendation(projects, 'refresh-token', clientFactory)).resolves.toEqual({ projectId: 'project-7' });
+        expect(vi.mocked(clientFactory)).toHaveBeenCalledTimes(8);
+        projects.forEach(({ projectId }, index) => {
+            const client = clients.get(projectId);
+            if (!client)
+                throw new Error('Unexpected project');
+            expect(vi.mocked(client.getDatasets)).toHaveBeenCalledTimes(index < 8 ? 1 : 0);
+            expect(vi.mocked(client.query)).toHaveBeenCalledTimes(index < 8 ? 3 : 0);
+        });
+    });
+});
 describe('BigqueryWarehouseClient.sanitizeLabelsWithValues', () => {
     let warnSpy;
     beforeEach(() => {

package/dist/warehouseClients/SnowflakeWarehouseClient.d.ts

@@ -49,6 +49,7 @@ export type SnowflakeSessionInventory = {
         name: string;
         comment: string | null;
         kind: string | null;
+        sizeBytes: number | null;
     }[];
     warehouses: {
         name: string;

package/dist/warehouseClients/SnowflakeWarehouseClient.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"SnowflakeWarehouseClient.d.ts","sourceRoot":"","sources":["../../src/warehouseClients/SnowflakeWarehouseClient.ts"],"names":[],"mappings":"AAAA,OAAO,EACH,OAAO,EACP,0BAA0B,EAC1B,aAAa,EAGb,MAAM,EAIN,mBAAmB,EAGnB,mBAAmB,EACnB,gBAAgB,EAChB,cAAc,EACd,KAAK,0BAA0B,EAC/B,KAAK,8BAA8B,EACtC,MAAM,mBAAmB,CAAC;AAE3B,OAAO,EAGH,UAAU,EACV,iBAAiB,EAEjB,cAAc,EAGjB,MAAM,eAAe,CAAC;AAGvB,OAAO,EAAE,gBAAgB,EAAE,MAAM,UAAU,CAAC;AAM5C,OAAO,mBAAmB,MAAM,uBAAuB,CAAC;AACxD,OAAO,uBAAuB,MAAM,2BAA2B,CAAC;AAchE,oBAAY,cAAc;IACtB,MAAM,WAAW;IACjB,OAAO,YAAY;IACnB,OAAO,YAAY;IACnB,OAAO,YAAY;IACnB,GAAG,QAAQ;IACX,MAAM,WAAW;IACjB,QAAQ,aAAa;IACrB,KAAK,UAAU;IACf,MAAM,WAAW;IACjB,MAAM,WAAW;IACjB,MAAM,WAAW;IACjB,gBAAgB,qBAAqB;IACrC,IAAI,SAAS;IACb,KAAK,UAAU;IACf,MAAM,WAAW;IACjB,IAAI,SAAS;IACb,OAAO,YAAY;IACnB,IAAI,SAAS;IACb,QAAQ,aAAa;IACrB,IAAI,SAAS;IACb,SAAS,cAAc;IACvB,aAAa,kBAAkB;IAC/B,aAAa,kBAAkB;IAC/B,YAAY,iBAAiB;IAC7B,OAAO,YAAY;IACnB,MAAM,WAAW;IACjB,KAAK,UAAU;IACf,SAAS,cAAc;CAC1B;AAqBD,MAAM,MAAM,oBAAoB,GAAG;IAC/B,WAAW,EAAE,MAAM,CAAC;IACpB,YAAY,EAAE,MAAM,CAAC;IACrB,SAAS,EAAE,IAAI,GAAG,IAAI,CAAC;CAC1B,CAAC;AAqFF,MAAM,MAAM,wBAAwB,GAAG;IACnC,IAAI,EAAE,MAAM,GAAG,IAAI,CAAC;IACpB,SAAS,EAAE,MAAM,GAAG,IAAI,CAAC;IACzB,QAAQ,EAAE,MAAM,GAAG,IAAI,CAAC;IACxB,MAAM,EAAE,MAAM,GAAG,IAAI,CAAC;CACzB,CAAC;AAEF,MAAM,MAAM,yBAAyB,GAAG;IACpC,SAAS,EAAE;QACP,IAAI,EAAE,MAAM,CAAC;QACb,OAAO,EAAE,MAAM,GAAG,IAAI,CAAC;QACvB,IAAI,EAAE,MAAM,GAAG,IAAI,CAAC;KACvB,EAAE,CAAC;IACJ,UAAU,EAAE;QACR,IAAI,EAAE,MAAM,CAAC;QACb,IAAI,EAAE,MAAM,GAAG,IAAI,CAAC;QACpB,KAAK,EAAE,MAAM,GAAG,IAAI,CAAC;QACrB,kBAAkB,EAAE,MAAM,GAAG,IAAI,CAAC;KACrC,EAAE,CAAC;IACJ,KAAK,EAAE;QACH,IAAI,EAAE,MAAM,CAAC;QACb,SAAS,EAAE,OAAO,CAAC;KACtB,EAAE,CAAC;IACJ,OAAO,EAAE;QACL,YAAY,EAAE,MAAM,CAAC;QACrB,IAAI,EAAE,MAAM,CAAC;KAChB,EAAE,CAAC;CACP,CAAC;AAEF,MAAM,MAAM,yBAAyB,GAAG;IACpC,IAAI,EAAE,MAAM,CAAC;IACb,QAAQ,EAAE,wBAAwB,CAAC;IACnC,SAAS,EAAE,yBAAyB,CAAC;CACxC,CAAC;AAEF,MAAM,MAAM,gCAAgC,GACtC,oBAAoB,GACpB,gBAAgB,GAChB,aAAa,GACb,iBAAiB,GACjB,kBAAkB,GAClB,gBAAgB,GAChB,SAAS,CAAC;AAEhB,MAAM,MAAM,+BAA+B,GAAG;IAC1C,QAAQ,EAAE,gCAAgC,CAAC;IAC3C,IAAI,EAAE,MAAM,GAAG,IAAI,CAAC;IACpB,gBAAgB,EAAE,MAAM,CAAC;CAC5B,CAAC;AAiGF,eAAO,MAAM,mBAAmB,GAAI,OAAO,MAAM,KAAG,MAGV,CAAC;AAE3C,eAAO,MAAM,2BAA2B,GACpC,OAAO,OAAO,KACf,+BA+DF,CAAC;AAEF,qBAAa,wBAAyB,SAAQ,KAAK;IAC/C,QAAQ,CAAC,OAAO,EAAE,+BAA+B,CAAC;IAElD,OAAO,CAAC,QAAQ,CAAC,QAAQ,CAAW;gBAExB,KAAK,EAAE,OAAO;IAa1B,WAAW,IAAI,OAAO;CAGzB;AAED,MAAM,MAAM,0BAA0B,GAAG;IACrC,UAAU,EAAE,UAAU,CAAC;CAC1B,CAAC;AAEF,MAAM,MAAM,sBAAsB,GAAG,gBAAgB,GAAG,kBAAkB,CAAC;AAE3E,MAAM,MAAM,uBAAuB,GAAG;IAClC,cAAc,EAAE,MAAM,GAAG,IAAI,CAAC;IAC9B,gBAAgB,EAAE,MAAM,GAAG,IAAI,CAAC;CACnC,CAAC;AAEF,eAAO,MAAM,YAAY,GAAI,MAAM,MAAM,KAAG,aA+B3C,CAAC;AAgBF,qBAAa,mBAAoB,SAAQ,uBAAuB;IAC5D,QAAQ,CAAC,IAAI,4BAA4B;IAEzC,cAAc,IAAI,mBAAmB;IAIrC,YAAY,CAAC,GAAG,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,MAAM;IAajD,YAAY,CAAC,KAAK,EAAE,MAAM,GAAG,MAAM;IAmBnC,uBAAuB,CACnB,iBAAiB,EAAE,MAAM,EACzB,eAAe,EAAE,MAAM,GACxB,MAAM;IAKT,UAAU,CAAC,QAAQ,EAAE,MAAM,EAAE,GAAG,MAAM;IAKtC,aAAa,CAAC,UAAU,EAAE,MAAM,EAAE,OAAO,CAAC,EAAE,MAAM,GAAG,MAAM;CAO9D;AAED,qBAAa,wBAAyB,SAAQ,mBAAmB,CAAC,0BAA0B,CAAC;IACzF,OAAO,CAAC,MAAM,CAAC,QAAQ,CAAC,oBAAoB,CAAQ;IAEpD,iBAAiB,EAAE,iBAAiB,CAAC;IAErC,2BAA2B,CAAC,EAAE,OAAO,CAAC;IAEtC,MAAM,CAAC,cAAc,CAAC,IAAI,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,GAAG,MAAM;IAO3D,OAAO,CAAC,4BAA4B,CAAC,CAAsB;IAE3D,OAAO,CAAC,QAAQ,CAAC,UAAU,CAAqB;IAEhD,OAAO,CAAC,QAAQ,CAAC,oBAAoB,CAAqB;IAE1D,OAAO,CAAC,QAAQ,CAAC,sBAAsB,CAAC,CAAkC;gBAE9D,WAAW,EAAE,0BAA0B;YAkHrC,aAAa;IAU3B,OAAO,CAAC,0BAA0B;YASpB,oBAAoB;YAwBpB,2BAA2B;YA+C3B,gBAAgB;IAoBxB,wBAAwB,IAAI,OAAO,CAAC,0BAA0B,CAAC;IAY/D,6BAA6B,CAC/B,OAAO,EAAE,0BAA0B,GACpC,OAAO,CAAC,IAAI,CAAC;IAQV,yBAAyB,CAC3B,OAAO,EAAE,0BAA0B,GACpC,OAAO,CAAC,IAAI,CAAC;IAOV,mBAAmB,CACrB,IAAI,CAAC,EAAE,MAAM,GACd,OAAO,CAAC,yBAAyB,CAAC;IA6E/B,qBAAqB,CACvB,IAAI,EAAE,MAAM,GACb,OAAO,CAAC,uBAAuB,CAAC;IA6B7B,gBAAgB,CAClB,IAAI,EAAE,MAAM,EACZ,IAAI,EAAE,sBAAsB,EAC5B,SAAS,EAAE,MAAM,GAClB,OAAO,CAAC,IAAI,CAAC;IAaV,kBAAkB,CACpB,IAAI,EAAE,MAAM,EACZ,IAAI,EAAE,sBAAsB,GAC7B,OAAO,CAAC,IAAI,CAAC;IAQhB,cAAc,IAAI,oBAAoB,GAAG,IAAI;IAIvC,6BAA6B,CAC/B,SAAS,GAAE,MAAsC,EACjD,YAAY,GAAE,MAAU,EACxB,yBAAyB,GAAE,MAAa,EACxC,eAAe,GAAE,MAAM,GAAG,IAAW,GACtC,OAAO,CAAC;QAAE,WAAW,EAAE,MAAM,CAAC;QAAC,SAAS,EAAE,MAAM,CAAA;KAAE,CAAC;YAoGxC,gBAAgB;IAmG9B,OAAO,CAAC,sBAAsB;YAkBhB,iBAAiB;IAuBzB,iBAAiB,CACnB,EAAE,GAAG,EAAE,MAAM,EAAE,IAAI,EAAE,QAAQ,EAAE,EAAE,8BAA8B,EAC/D,qBAAqB,CAAC,EAAE,CACpB,IAAI,EAAE,gBAAgB,CAAC,MAAM,CAAC,EAC9B,MAAM,EAAE,gBAAgB,CAAC,QAAQ,CAAC,KACjC,IAAI,GAAG,OAAO,CAAC,IAAI,CAAC,GAC1B,OAAO,CAAC,0BAA0B,CAAC;YA2CxB,qBAAqB;IA2I7B,WAAW,CACb,GAAG,EAAE,MAAM,EACX,cAAc,EAAE,CAAC,IAAI,EAAE,gBAAgB,KAAK,IAAI,GAAG,OAAO,CAAC,IAAI,CAAC,EAChE,OAAO,EAAE;QACL,MAAM,CAAC,EAAE,OAAO,EAAE,CAAC;QACnB,IAAI,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;QAC9B,QAAQ,CAAC,EAAE,MAAM,CAAC;KACrB,GACF,OAAO,CAAC,IAAI,CAAC;YAuBF,sBAAsB;YA6DtB,iBAAiB;YAmCjB,oBAAoB;IA0B5B,UAAU,CACZ,MAAM,EAAE;QACJ,QAAQ,EAAE,MAAM,CAAC;QACjB,MAAM,EAAE,MAAM,CAAC;QACf,KAAK,EAAE,MAAM,CAAC;KACjB,EAAE;IAuCD,YAAY;;;;;IA8BZ,SAAS,CACX,SAAS,EAAE,MAAM,EACjB,MAAM,CAAC,EAAE,MAAM,EACf,QAAQ,CAAC,EAAE,MAAM,EACjB,IAAI,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,GAC9B,OAAO,CAAC,gBAAgB,CAAC;IAsC5B,OAAO,CAAC,wBAAwB;IA0ChC,OAAO,CAAC,2BAA2B;IAcnC,UAAU,CAAC,KAAK,EAAE,cAAc,EAAE,KAAK,GAAE,MAAW;CA8FvD"}
+{"version":3,"file":"SnowflakeWarehouseClient.d.ts","sourceRoot":"","sources":["../../src/warehouseClients/SnowflakeWarehouseClient.ts"],"names":[],"mappings":"AAAA,OAAO,EACH,OAAO,EACP,0BAA0B,EAC1B,aAAa,EAGb,MAAM,EAIN,mBAAmB,EAGnB,mBAAmB,EACnB,gBAAgB,EAChB,cAAc,EACd,KAAK,0BAA0B,EAC/B,KAAK,8BAA8B,EACtC,MAAM,mBAAmB,CAAC;AAE3B,OAAO,EAGH,UAAU,EACV,iBAAiB,EAEjB,cAAc,EAGjB,MAAM,eAAe,CAAC;AAGvB,OAAO,EAAE,gBAAgB,EAAE,MAAM,UAAU,CAAC;AAM5C,OAAO,mBAAmB,MAAM,uBAAuB,CAAC;AACxD,OAAO,uBAAuB,MAAM,2BAA2B,CAAC;AAchE,oBAAY,cAAc;IACtB,MAAM,WAAW;IACjB,OAAO,YAAY;IACnB,OAAO,YAAY;IACnB,OAAO,YAAY;IACnB,GAAG,QAAQ;IACX,MAAM,WAAW;IACjB,QAAQ,aAAa;IACrB,KAAK,UAAU;IACf,MAAM,WAAW;IACjB,MAAM,WAAW;IACjB,MAAM,WAAW;IACjB,gBAAgB,qBAAqB;IACrC,IAAI,SAAS;IACb,KAAK,UAAU;IACf,MAAM,WAAW;IACjB,IAAI,SAAS;IACb,OAAO,YAAY;IACnB,IAAI,SAAS;IACb,QAAQ,aAAa;IACrB,IAAI,SAAS;IACb,SAAS,cAAc;IACvB,aAAa,kBAAkB;IAC/B,aAAa,kBAAkB;IAC/B,YAAY,iBAAiB;IAC7B,OAAO,YAAY;IACnB,MAAM,WAAW;IACjB,KAAK,UAAU;IACf,SAAS,cAAc;CAC1B;AAsBD,MAAM,MAAM,oBAAoB,GAAG;IAC/B,WAAW,EAAE,MAAM,CAAC;IACpB,YAAY,EAAE,MAAM,CAAC;IACrB,SAAS,EAAE,IAAI,GAAG,IAAI,CAAC;CAC1B,CAAC;AAqFF,MAAM,MAAM,wBAAwB,GAAG;IACnC,IAAI,EAAE,MAAM,GAAG,IAAI,CAAC;IACpB,SAAS,EAAE,MAAM,GAAG,IAAI,CAAC;IACzB,QAAQ,EAAE,MAAM,GAAG,IAAI,CAAC;IACxB,MAAM,EAAE,MAAM,GAAG,IAAI,CAAC;CACzB,CAAC;AAEF,MAAM,MAAM,yBAAyB,GAAG;IACpC,SAAS,EAAE;QACP,IAAI,EAAE,MAAM,CAAC;QACb,OAAO,EAAE,MAAM,GAAG,IAAI,CAAC;QACvB,IAAI,EAAE,MAAM,GAAG,IAAI,CAAC;QACpB,SAAS,EAAE,MAAM,GAAG,IAAI,CAAC;KAC5B,EAAE,CAAC;IACJ,UAAU,EAAE;QACR,IAAI,EAAE,MAAM,CAAC;QACb,IAAI,EAAE,MAAM,GAAG,IAAI,CAAC;QACpB,KAAK,EAAE,MAAM,GAAG,IAAI,CAAC;QACrB,kBAAkB,EAAE,MAAM,GAAG,IAAI,CAAC;KACrC,EAAE,CAAC;IACJ,KAAK,EAAE;QACH,IAAI,EAAE,MAAM,CAAC;QACb,SAAS,EAAE,OAAO,CAAC;KACtB,EAAE,CAAC;IACJ,OAAO,EAAE;QACL,YAAY,EAAE,MAAM,CAAC;QACrB,IAAI,EAAE,MAAM,CAAC;KAChB,EAAE,CAAC;CACP,CAAC;AAEF,MAAM,MAAM,yBAAyB,GAAG;IACpC,IAAI,EAAE,MAAM,CAAC;IACb,QAAQ,EAAE,wBAAwB,CAAC;IACnC,SAAS,EAAE,yBAAyB,CAAC;CACxC,CAAC;AAEF,MAAM,MAAM,gCAAgC,GACtC,oBAAoB,GACpB,gBAAgB,GAChB,aAAa,GACb,iBAAiB,GACjB,kBAAkB,GAClB,gBAAgB,GAChB,SAAS,CAAC;AAEhB,MAAM,MAAM,+BAA+B,GAAG;IAC1C,QAAQ,EAAE,gCAAgC,CAAC;IAC3C,IAAI,EAAE,MAAM,GAAG,IAAI,CAAC;IACpB,gBAAgB,EAAE,MAAM,CAAC;CAC5B,CAAC;AAqJF,eAAO,MAAM,mBAAmB,GAAI,OAAO,MAAM,KAAG,MAGV,CAAC;AAE3C,eAAO,MAAM,2BAA2B,GACpC,OAAO,OAAO,KACf,+BA+DF,CAAC;AAEF,qBAAa,wBAAyB,SAAQ,KAAK;IAC/C,QAAQ,CAAC,OAAO,EAAE,+BAA+B,CAAC;IAElD,OAAO,CAAC,QAAQ,CAAC,QAAQ,CAAW;gBAExB,KAAK,EAAE,OAAO;IAa1B,WAAW,IAAI,OAAO;CAGzB;AAED,MAAM,MAAM,0BAA0B,GAAG;IACrC,UAAU,EAAE,UAAU,CAAC;CAC1B,CAAC;AAEF,MAAM,MAAM,sBAAsB,GAAG,gBAAgB,GAAG,kBAAkB,CAAC;AAE3E,MAAM,MAAM,uBAAuB,GAAG;IAClC,cAAc,EAAE,MAAM,GAAG,IAAI,CAAC;IAC9B,gBAAgB,EAAE,MAAM,GAAG,IAAI,CAAC;CACnC,CAAC;AAEF,eAAO,MAAM,YAAY,GAAI,MAAM,MAAM,KAAG,aA+B3C,CAAC;AAgBF,qBAAa,mBAAoB,SAAQ,uBAAuB;IAC5D,QAAQ,CAAC,IAAI,4BAA4B;IAEzC,cAAc,IAAI,mBAAmB;IAIrC,YAAY,CAAC,GAAG,EAAE,MAAM,EAAE,MAAM,EAAE,MAAM,GAAG,MAAM;IAajD,YAAY,CAAC,KAAK,EAAE,MAAM,GAAG,MAAM;IAmBnC,uBAAuB,CACnB,iBAAiB,EAAE,MAAM,EACzB,eAAe,EAAE,MAAM,GACxB,MAAM;IAKT,UAAU,CAAC,QAAQ,EAAE,MAAM,EAAE,GAAG,MAAM;IAKtC,aAAa,CAAC,UAAU,EAAE,MAAM,EAAE,OAAO,CAAC,EAAE,MAAM,GAAG,MAAM;CAO9D;AAED,qBAAa,wBAAyB,SAAQ,mBAAmB,CAAC,0BAA0B,CAAC;IACzF,OAAO,CAAC,MAAM,CAAC,QAAQ,CAAC,oBAAoB,CAAQ;IAEpD,iBAAiB,EAAE,iBAAiB,CAAC;IAErC,2BAA2B,CAAC,EAAE,OAAO,CAAC;IAEtC,MAAM,CAAC,cAAc,CAAC,IAAI,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,GAAG,MAAM;IAO3D,OAAO,CAAC,4BAA4B,CAAC,CAAsB;IAE3D,OAAO,CAAC,QAAQ,CAAC,UAAU,CAAqB;IAEhD,OAAO,CAAC,QAAQ,CAAC,oBAAoB,CAAqB;IAE1D,OAAO,CAAC,QAAQ,CAAC,sBAAsB,CAAC,CAAkC;gBAE9D,WAAW,EAAE,0BAA0B;YAkHrC,aAAa;IAU3B,OAAO,CAAC,0BAA0B;YASpB,oBAAoB;YAwBpB,2BAA2B;YA+C3B,gBAAgB;IAoBxB,wBAAwB,IAAI,OAAO,CAAC,0BAA0B,CAAC;IAY/D,6BAA6B,CAC/B,OAAO,EAAE,0BAA0B,GACpC,OAAO,CAAC,IAAI,CAAC;IAQV,yBAAyB,CAC3B,OAAO,EAAE,0BAA0B,GACpC,OAAO,CAAC,IAAI,CAAC;IAOV,mBAAmB,CACrB,IAAI,CAAC,EAAE,MAAM,GACd,OAAO,CAAC,yBAAyB,CAAC;IAwI/B,qBAAqB,CACvB,IAAI,EAAE,MAAM,GACb,OAAO,CAAC,uBAAuB,CAAC;IA6B7B,gBAAgB,CAClB,IAAI,EAAE,MAAM,EACZ,IAAI,EAAE,sBAAsB,EAC5B,SAAS,EAAE,MAAM,GAClB,OAAO,CAAC,IAAI,CAAC;IAaV,kBAAkB,CACpB,IAAI,EAAE,MAAM,EACZ,IAAI,EAAE,sBAAsB,GAC7B,OAAO,CAAC,IAAI,CAAC;IAQhB,cAAc,IAAI,oBAAoB,GAAG,IAAI;IAIvC,6BAA6B,CAC/B,SAAS,GAAE,MAAsC,EACjD,YAAY,GAAE,MAAU,EACxB,yBAAyB,GAAE,MAAa,EACxC,eAAe,GAAE,MAAM,GAAG,IAAW,GACtC,OAAO,CAAC;QAAE,WAAW,EAAE,MAAM,CAAC;QAAC,SAAS,EAAE,MAAM,CAAA;KAAE,CAAC;YAoGxC,gBAAgB;IAmG9B,OAAO,CAAC,sBAAsB;YAkBhB,iBAAiB;IAuBzB,iBAAiB,CACnB,EAAE,GAAG,EAAE,MAAM,EAAE,IAAI,EAAE,QAAQ,EAAE,EAAE,8BAA8B,EAC/D,qBAAqB,CAAC,EAAE,CACpB,IAAI,EAAE,gBAAgB,CAAC,MAAM,CAAC,EAC9B,MAAM,EAAE,gBAAgB,CAAC,QAAQ,CAAC,KACjC,IAAI,GAAG,OAAO,CAAC,IAAI,CAAC,GAC1B,OAAO,CAAC,0BAA0B,CAAC;YA2CxB,qBAAqB;IA2I7B,WAAW,CACb,GAAG,EAAE,MAAM,EACX,cAAc,EAAE,CAAC,IAAI,EAAE,gBAAgB,KAAK,IAAI,GAAG,OAAO,CAAC,IAAI,CAAC,EAChE,OAAO,EAAE;QACL,MAAM,CAAC,EAAE,OAAO,EAAE,CAAC;QACnB,IAAI,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,CAAC;QAC9B,QAAQ,CAAC,EAAE,MAAM,CAAC;KACrB,GACF,OAAO,CAAC,IAAI,CAAC;YAuBF,sBAAsB;YA6DtB,iBAAiB;YAmCjB,oBAAoB;IA0B5B,UAAU,CACZ,MAAM,EAAE;QACJ,QAAQ,EAAE,MAAM,CAAC;QACjB,MAAM,EAAE,MAAM,CAAC;QACf,KAAK,EAAE,MAAM,CAAC;KACjB,EAAE;IAuCD,YAAY;;;;;IA8BZ,SAAS,CACX,SAAS,EAAE,MAAM,EACjB,MAAM,CAAC,EAAE,MAAM,EACf,QAAQ,CAAC,EAAE,MAAM,EACjB,IAAI,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,MAAM,CAAC,GAC9B,OAAO,CAAC,gBAAgB,CAAC;IAsC5B,OAAO,CAAC,wBAAwB;IA0ChC,OAAO,CAAC,2BAA2B;IAcnC,UAAU,CAAC,KAAK,EAAE,cAAc,EAAE,KAAK,GAAE,MAAW;CA8FvD"}

package/dist/warehouseClients/SnowflakeWarehouseClient.js

@@ -96,6 +96,7 @@ const EXTERNAL_BROWSER_AUTHENTICATOR = 'EXTERNALBROWSER';
 const OAUTH_AUTHORIZATION_CODE_AUTHENTICATOR = 'OAUTH_AUTHORIZATION_CODE';
 const SESSION_DISCOVERY_LIMIT = 100;
 const SESSION_SCHEMA_DISCOVERY_LIMIT = 1000;
+const SESSION_DATABASE_SIZE_LIMIT = 25;
 // Minimum allowed by Snowflake (range 16-160, default 160)
 const SNOWFLAKE_RESULT_CHUNK_SIZE_MB = 16;
 const getAccessTokenExpiresAt = (token) => {
@@ -202,7 +203,7 @@ const listSnowflakeSchemas = (rows) => rows.flatMap((rawRow) => {
     }
     return [{ databaseName, name }];
 });
-const listSnowflakeDatabases = (rows) => rows
+const listSnowflakeDatabases = (rows, databaseSizes) => rows
     .flatMap((rawRow) => {
     const row = rawRow;
     const name = getSnowflakeRowValue(row, 'name');
@@ -216,6 +217,7 @@ const listSnowflakeDatabases = (rows) => rows
                 ? null
                 : getSnowflakeRowValue(row, 'comment'),
             kind: getSnowflakeRowValue(row, 'kind'),
+            sizeBytes: databaseSizes.get(name) ?? null,
         },
     ];
 })
@@ -237,6 +239,44 @@ const listSnowflakeWarehouses = (rows) => rows
     ];
 })
     .slice(0, SESSION_DISCOVERY_LIMIT);
+const WAREHOUSE_SIZE_ORDER = [
+    'X-Small',
+    'Small',
+    'Medium',
+    'Large',
+    'X-Large',
+    '2X-Large',
+    '3X-Large',
+    '4X-Large',
+    '5X-Large',
+    '6X-Large',
+];
+const warehouseSizeRank = (size) => {
+    const index = WAREHOUSE_SIZE_ORDER.findIndex((candidate) => candidate.toLowerCase() === (size ?? '').trim().toLowerCase());
+    return index === -1 ? Number.MAX_SAFE_INTEGER : index;
+};
+const pickSizeDiscoveryWarehouse = (rows) => {
+    const warehouses = rows
+        .map((rawRow) => {
+        const row = rawRow;
+        return {
+            name: getSnowflakeRowValue(row, 'name'),
+            size: getSnowflakeRowValue(row, 'size'),
+            state: (getSnowflakeRowValue(row, 'state') ?? '')
+                .trim()
+                .toUpperCase(),
+        };
+    })
+        .filter((warehouse) => Boolean(warehouse.name));
+    const started = warehouses.filter((warehouse) => warehouse.state === 'STARTED');
+    const pool = started.length > 0 ? started : warehouses;
+    const best = pool.reduce((currentBest, candidate) => currentBest === null ||
+        warehouseSizeRank(candidate.size) <
+            warehouseSizeRank(currentBest.size)
+        ? candidate
+        : currentBest, null);
+    return best?.name ?? null;
+};
 const snowflakeIdentifier = (value) => /^[A-Za-z_][A-Za-z0-9_$]*$/.test(value)
     ? value
     : `"${value.replace(/"/g, '""')}"`;
@@ -623,6 +663,38 @@ class SnowflakeWarehouseClient extends WarehouseBaseClient_1.default {
         catch {
             schemaRows = [];
         }
+        const databaseSizes = new Map();
+        const sizeCandidates = databasesResult.rows
+            .map((rawRow) => getSnowflakeRowValue(rawRow, 'name'))
+            .filter((name) => Boolean(name))
+            .slice(0, SESSION_DATABASE_SIZE_LIMIT);
+        // INFORMATION_SCHEMA aggregation needs an active warehouse; without a
+        // session default, activate one (prefer already-running, then smallest)
+        if (sizeCandidates.length > 0) {
+            const sessionWarehouseName = getSnowflakeRowValue(defaultsRow, 'warehouse');
+            const sizeWarehouseName = sessionWarehouseName ??
+                pickSizeDiscoveryWarehouse(warehousesResult.rows);
+            if (!sessionWarehouseName && sizeWarehouseName) {
+                try {
+                    await this.executeStatements(connection, `USE WAREHOUSE ${(0, exports.snowflakeIdentifier)(sizeWarehouseName)}`);
+                }
+                catch {
+                    /* size queries will fail and sizes stay null */
+                }
+            }
+        }
+        await sizeCandidates.reduce((previous, databaseName) => previous.then(async () => {
+            try {
+                const sizeResult = await this.executeStatements(connection, `SELECT SUM(bytes) AS "total_bytes" FROM ${(0, exports.snowflakeIdentifier)(databaseName)}.INFORMATION_SCHEMA.TABLES WHERE table_type = 'BASE TABLE'`);
+                const totalBytes = getSnowflakeRowNumber((sizeResult.rows[0] ?? {}), 'total_bytes');
+                if (totalBytes !== null && totalBytes > 0) {
+                    databaseSizes.set(databaseName, totalBytes);
+                }
+            }
+            catch {
+                /* size stays null for this database */
+            }
+        }), Promise.resolve());
         const defaultRole = getSnowflakeRowValue(defaultsRow, 'role');
         const grantedRoles = [
             ...new Set([
@@ -649,7 +721,7 @@ class SnowflakeWarehouseClient extends WarehouseBaseClient_1.default {
                 schema: getSnowflakeRowValue(defaultsRow, 'schema'),
             },
             inventory: {
-                databases: listSnowflakeDatabases(databasesResult.rows),
+                databases: listSnowflakeDatabases(databasesResult.rows, databaseSizes),
                 warehouses: listSnowflakeWarehouses(warehousesResult.rows),
                 roles,
                 schemas: listSnowflakeSchemas(schemaRows),

package/dist/warehouseClients/SnowflakeWarehouseClient.test.js

@@ -311,6 +311,17 @@ describe('SnowflakeWarehouseClient', () => {
                     },
                 ];
             }
+            else if (sqlText.startsWith('SELECT SUM(bytes)')) {
+                if (sqlText.includes('DATABASE_0.')) {
+                    rows = [{ total_bytes: 1250 }];
+                }
+                else if (sqlText.includes('DATABASE_1.')) {
+                    rows = [{ total_bytes: 500 }];
+                }
+                else {
+                    rows = [{ total_bytes: null }];
+                }
+            }
             complete(undefined, { getColumns: () => SnowflakeWarehouseClient_mock_1.queryColumnsMock }, rows);
         });
         vi.mocked(snowflake_sdk_1.createConnection).mockImplementationOnce(() => interactiveConnectionMock(execute));
@@ -331,6 +342,7 @@ describe('SnowflakeWarehouseClient', () => {
                     name: `DATABASE_${index}`,
                     comment: index === 0 ? 'Primary analytics' : null,
                     kind: index === 1 ? 'IMPORTED DATABASE' : 'STANDARD',
+                    sizeBytes: [1250, 500][index] ?? null,
                 })),
                 warehouses: [
                     {
@@ -354,10 +366,56 @@ describe('SnowflakeWarehouseClient', () => {
         expect(execute).toHaveBeenCalledWith(expect.objectContaining({
             sqlText: 'SHOW SCHEMAS IN ACCOUNT LIMIT 1000',
         }));
+        expect(execute).toHaveBeenCalledWith(expect.objectContaining({
+            sqlText: `SELECT SUM(bytes) AS "total_bytes" FROM DATABASE_0.INFORMATION_SCHEMA.TABLES WHERE table_type = 'BASE TABLE'`,
+        }));
+        expect(execute.mock.calls.filter(([statement]) => statement.sqlText.startsWith('SELECT SUM(bytes)'))).toHaveLength(25);
+        expect(execute.mock.calls.some(([statement]) => statement.sqlText.startsWith('USE WAREHOUSE'))).toBe(false);
         expect(execute).toHaveBeenCalledWith(expect.objectContaining({
             sqlText: 'SHOW GRANTS TO USER "user.name@example.com" LIMIT 100',
         }));
     });
+    it('returns null database sizes when size discovery fails', async () => {
+        const execute = vi.fn(({ sqlText, complete, }) => {
+            if (sqlText.startsWith('SELECT SUM(bytes)')) {
+                complete(new Error('Insufficient privileges'), { getColumns: () => SnowflakeWarehouseClient_mock_1.queryColumnsMock }, []);
+                return;
+            }
+            let rows = [];
+            if (sqlText.startsWith('SELECT CURRENT_USER')) {
+                rows = [{ USER: 'test-user' }];
+            }
+            else if (sqlText.startsWith('SHOW DATABASES')) {
+                rows = [{ name: 'ANALYTICS' }, { name: 'RAW' }];
+            }
+            else if (sqlText.startsWith('SHOW WAREHOUSES')) {
+                rows = [
+                    { name: 'WH_BIG', size: 'Large', state: 'SUSPENDED' },
+                    {
+                        name: 'WH_SMALL',
+                        size: 'X-Small',
+                        state: 'STARTED',
+                    },
+                ];
+            }
+            complete(undefined, { getColumns: () => SnowflakeWarehouseClient_mock_1.queryColumnsMock }, rows);
+        });
+        vi.mocked(snowflake_sdk_1.createConnection).mockImplementationOnce(() => interactiveConnectionMock(execute));
+        const warehouse = new SnowflakeWarehouseClient_1.SnowflakeWarehouseClient(SnowflakeWarehouseClient_mock_1.credentials);
+        const discovery = await warehouse.getSessionDiscovery('test-user');
+        expect(discovery.inventory.databases).toEqual([
+            {
+                name: 'ANALYTICS',
+                comment: null,
+                kind: null,
+                sizeBytes: null,
+            },
+            { name: 'RAW', comment: null, kind: null, sizeBytes: null },
+        ]);
+        expect(execute).toHaveBeenCalledWith(expect.objectContaining({
+            sqlText: 'USE WAREHOUSE WH_SMALL',
+        }));
+    });
     it('opens, checks, and closes a diagnostic connection', async () => {
         const execute = vi.fn(({ sqlText, complete, }) => complete(undefined, { getColumns: () => SnowflakeWarehouseClient_mock_1.queryColumnsMock }, []));
         const connection = interactiveConnectionMock(execute);

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@lightdash/warehouses",
-  "version": "0.3414.0",
+  "version": "0.3416.0",
   "description": "Warehouse connectors for Lightdash",
   "license": "MIT",
   "repository": {
@@ -30,7 +30,7 @@
     "snowflake-sdk": "2.3.4",
     "ssh2": "1.14.0",
     "trino-client": "0.2.9",
-    "@lightdash/common": "0.3414.0"
+    "@lightdash/common": "0.3416.0"
   },
   "devDependencies": {
     "@types/big.js": "6.2.2",