@lightdash/warehouses 0.3413.1 → 0.3415.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
- import { BigQuery, Dataset, QueryRowsResponse } from '@google-cloud/bigquery';
1
+ import { BigQuery, Dataset, DatasetsResponse, GetDatasetsOptions, QueryRowsResponse } from '@google-cloud/bigquery';
2
2
  import bigquery from '@google-cloud/bigquery/build/src/types';
3
- import { AnyType, BigqueryDataset, BigqueryProject, CreateBigqueryCredentials, DimensionType, Metric, PartitionColumn, SupportedDbtAdapter, TimeIntervalUnit, WarehouseQueryError, WarehouseResults, WarehouseTypes } from '@lightdash/common';
3
+ import { AnyType, BigqueryDataset, BigqueryProject, BigqueryProjectRecommendation, CreateBigqueryCredentials, DimensionType, Metric, PartitionColumn, SupportedDbtAdapter, TimeIntervalUnit, WarehouseQueryError, WarehouseResults, WarehouseTypes } from '@lightdash/common';
4
4
  import { WarehouseCatalog, WarehouseExecuteAsyncQuery, WarehouseExecuteAsyncQueryArgs } from '../types';
5
5
  import WarehouseBaseClient from './WarehouseBaseClient';
6
6
  import WarehouseBaseSqlBuilder from './WarehouseBaseSqlBuilder';
@@ -90,7 +90,17 @@ export declare class BigqueryWarehouseClient extends WarehouseBaseClient<CreateB
90
90
  parseError(error: bigquery.IErrorProto, query?: string): WarehouseQueryError;
91
91
  executeAsyncQuery({ sql, tags, timezone }: WarehouseExecuteAsyncQueryArgs, resultsStreamCallback?: (rows: WarehouseResults['rows'], fields: WarehouseResults['fields']) => void): Promise<WarehouseExecuteAsyncQuery>;
92
92
  private awaitJobCompletion;
93
- static getDatabases(projectId: string, refresh_token: string): Promise<BigqueryDataset[]>;
93
+ static getDatabases(projectId: string, refresh_token: string, bigqueryClient?: {
94
+ getDatasets: () => Promise<DatasetsResponse>;
95
+ query: (query: string) => Promise<QueryRowsResponse>;
96
+ }): Promise<BigqueryDataset[]>;
97
+ static getProjectRecommendation(projects: BigqueryProject[], refreshToken: string, bigqueryClientFactory?: (projectId: string) => {
98
+ getDatasets: (options: GetDatasetsOptions) => Promise<DatasetsResponse>;
99
+ query: (options: {
100
+ query: string;
101
+ location: string;
102
+ }) => Promise<QueryRowsResponse>;
103
+ }): Promise<BigqueryProjectRecommendation>;
94
104
  static getProjects(accessToken: string): Promise<BigqueryProject[]>;
95
105
  }
96
106
  export {};
@@ -1 +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"}
1
+ {"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"}
@@ -12,6 +12,8 @@ const processPromisesInBatches_1 = require("../utils/processPromisesInBatches");
12
12
  const sql_1 = require("../utils/sql");
13
13
  const WarehouseBaseClient_1 = __importDefault(require("./WarehouseBaseClient"));
14
14
  const WarehouseBaseSqlBuilder_1 = __importDefault(require("./WarehouseBaseSqlBuilder"));
15
+ const MAX_RECOMMENDATION_PROJECTS = 8;
16
+ const MAX_RECOMMENDATION_REGIONS_PER_PROJECT = 3;
15
17
  var BigqueryFieldType;
16
18
  (function (BigqueryFieldType) {
17
19
  BigqueryFieldType["STRING"] = "STRING";
@@ -516,23 +518,112 @@ class BigqueryWarehouseClient extends WarehouseBaseClient_1.default {
516
518
  });
517
519
  });
518
520
  }
519
- static async getDatabases(projectId, refresh_token) {
520
- const bigqueryClient = new bigquery_1.BigQuery({
521
+ static async getDatabases(projectId, refresh_token, bigqueryClient = new bigquery_1.BigQuery({
522
+ projectId,
523
+ credentials: {
524
+ type: 'authorized_user',
525
+ client_id: process.env.AUTH_GOOGLE_OAUTH2_CLIENT_ID,
526
+ client_secret: process.env.AUTH_GOOGLE_OAUTH2_CLIENT_SECRET,
527
+ refresh_token,
528
+ },
529
+ })) {
530
+ const datasets = await bigqueryClient.getDatasets();
531
+ const databases = await Promise.all(datasets[0].map(async (dataset, index) => {
532
+ let sizeBytes = null;
533
+ if (index < 25) {
534
+ try {
535
+ const [rows] = await bigqueryClient.query(`SELECT SUM(size_bytes) AS sizeBytes FROM \`${dataset.projectId}.${dataset.id}.__TABLES__\``);
536
+ const size = Number(rows[0]?.sizeBytes);
537
+ sizeBytes = Number.isFinite(size) ? size : null;
538
+ }
539
+ catch {
540
+ sizeBytes = null;
541
+ }
542
+ }
543
+ return {
544
+ projectId: dataset.projectId,
545
+ location: dataset.location,
546
+ datasetId: dataset.id,
547
+ sizeBytes,
548
+ };
549
+ }));
550
+ return databases;
551
+ }
552
+ static async getProjectRecommendation(projects, refreshToken, bigqueryClientFactory = (projectId) => {
553
+ const client = new bigquery_1.BigQuery({
521
554
  projectId,
522
555
  credentials: {
523
556
  type: 'authorized_user',
524
557
  client_id: process.env.AUTH_GOOGLE_OAUTH2_CLIENT_ID,
525
558
  client_secret: process.env.AUTH_GOOGLE_OAUTH2_CLIENT_SECRET,
526
- refresh_token,
559
+ refresh_token: refreshToken,
527
560
  },
528
561
  });
529
- const datasets = await bigqueryClient.getDatasets();
530
- const databases = datasets[0].map((d) => ({
531
- projectId: d.projectId,
532
- location: d.location,
533
- datasetId: d.id,
562
+ return {
563
+ getDatasets: (options) => client.getDatasets(options),
564
+ query: (options) => client.query(options.query, { location: options.location }),
565
+ };
566
+ }) {
567
+ const projectScores = await Promise.all(projects
568
+ .slice(0, MAX_RECOMMENDATION_PROJECTS)
569
+ .map(async ({ projectId }) => {
570
+ try {
571
+ const client = bigqueryClientFactory(projectId);
572
+ const [datasets] = await client.getDatasets({
573
+ autoPaginate: false,
574
+ maxResults: 1000,
575
+ });
576
+ const locationsByQualifier = new Map();
577
+ datasets.forEach((dataset) => {
578
+ if (dataset.location) {
579
+ locationsByQualifier.set(dataset.location.toLowerCase(), dataset.location);
580
+ }
581
+ });
582
+ const locations = Array.from(locationsByQualifier.entries()).slice(0, MAX_RECOMMENDATION_REGIONS_PER_PROJECT);
583
+ const regionScores = await Promise.all(locations.map(async ([qualifier, location]) => {
584
+ try {
585
+ const [rows] = await client.query({
586
+ query: `SELECT table_schema, SUM(total_logical_bytes) AS total_bytes FROM \`${projectId}\`.\`region-${qualifier}\`.INFORMATION_SCHEMA.TABLE_STORAGE GROUP BY table_schema`,
587
+ location,
588
+ });
589
+ let largestDatasetSize = null;
590
+ rows.forEach((row) => {
591
+ const size = Number(row.total_bytes);
592
+ if (Number.isFinite(size) &&
593
+ (largestDatasetSize === null ||
594
+ size > largestDatasetSize)) {
595
+ largestDatasetSize = size;
596
+ }
597
+ });
598
+ return {
599
+ error: false,
600
+ sizeBytes: largestDatasetSize,
601
+ };
602
+ }
603
+ catch {
604
+ return { error: true, sizeBytes: null };
605
+ }
606
+ }));
607
+ if (regionScores.some(({ error }) => error)) {
608
+ return { projectId, sizeBytes: null };
609
+ }
610
+ const sizeBytes = regionScores.reduce((largestSize, regionScore) => regionScore.sizeBytes !== null &&
611
+ (largestSize === null ||
612
+ regionScore.sizeBytes > largestSize)
613
+ ? regionScore.sizeBytes
614
+ : largestSize, null);
615
+ return { projectId, sizeBytes };
616
+ }
617
+ catch {
618
+ return { projectId, sizeBytes: null };
619
+ }
534
620
  }));
535
- return databases;
621
+ const recommendation = projectScores.reduce((largestProject, projectScore) => projectScore.sizeBytes !== null &&
622
+ (largestProject.sizeBytes === null ||
623
+ projectScore.sizeBytes > largestProject.sizeBytes)
624
+ ? projectScore
625
+ : largestProject, { projectId: null, sizeBytes: null });
626
+ return { projectId: recommendation.projectId };
536
627
  }
537
628
  static async getProjects(accessToken) {
538
629
  // Use BigQuery REST API directly instead of the SDK because:
@@ -42,6 +42,172 @@ describe('BigqueryWarehouseClient', () => {
42
42
  expect(BigqueryWarehouseClient_mock_1.getTableResponse.getMetadata).toHaveBeenCalledTimes(1);
43
43
  });
44
44
  });
45
+ describe('BigqueryWarehouseClient.getDatabases', () => {
46
+ const createDataset = (datasetId) => new bigquery_1.Dataset(new bigquery_1.BigQuery({ projectId: 'test-project' }), datasetId, {
47
+ location: 'EU',
48
+ });
49
+ const createClient = (datasets, query) => ({
50
+ getDatasets: vi
51
+ .fn()
52
+ .mockResolvedValue([datasets]),
53
+ query: vi.fn(query),
54
+ });
55
+ it('includes the stored size for each dataset', async () => {
56
+ const datasets = [createDataset('small'), createDataset('large')];
57
+ const queryMock = vi
58
+ .fn()
59
+ .mockResolvedValueOnce([[{ sizeBytes: 100 }]])
60
+ .mockResolvedValueOnce([[{ sizeBytes: 500 }]]);
61
+ const client = createClient(datasets, queryMock);
62
+ await expect(BigqueryWarehouseClient_1.BigqueryWarehouseClient.getDatabases('test-project', 'refresh-token', client)).resolves.toEqual([
63
+ {
64
+ projectId: 'test-project',
65
+ datasetId: 'small',
66
+ location: 'EU',
67
+ sizeBytes: 100,
68
+ },
69
+ {
70
+ projectId: 'test-project',
71
+ datasetId: 'large',
72
+ location: 'EU',
73
+ sizeBytes: 500,
74
+ },
75
+ ]);
76
+ expect(queryMock).toHaveBeenNthCalledWith(1, 'SELECT SUM(size_bytes) AS sizeBytes FROM `test-project.small.__TABLES__`');
77
+ });
78
+ it('returns null when an individual dataset size query fails', async () => {
79
+ const datasets = [createDataset('unavailable'), createDataset('ok')];
80
+ const queryMock = vi
81
+ .fn()
82
+ .mockRejectedValueOnce(new Error('permission denied'))
83
+ .mockResolvedValueOnce([[{ sizeBytes: 500 }]]);
84
+ const client = createClient(datasets, queryMock);
85
+ const result = await BigqueryWarehouseClient_1.BigqueryWarehouseClient.getDatabases('test-project', 'refresh-token', client);
86
+ expect(result.map(({ datasetId, sizeBytes }) => ({
87
+ datasetId,
88
+ sizeBytes,
89
+ }))).toEqual([
90
+ { datasetId: 'unavailable', sizeBytes: null },
91
+ { datasetId: 'ok', sizeBytes: 500 },
92
+ ]);
93
+ });
94
+ it('only queries sizes for the first 25 datasets', async () => {
95
+ const datasets = Array.from({ length: 30 }, (_, index) => createDataset(`dataset_${index}`));
96
+ const queryMock = vi
97
+ .fn()
98
+ .mockResolvedValue([[{ sizeBytes: 100 }]]);
99
+ const client = createClient(datasets, queryMock);
100
+ const result = await BigqueryWarehouseClient_1.BigqueryWarehouseClient.getDatabases('test-project', 'refresh-token', client);
101
+ expect(queryMock).toHaveBeenCalledTimes(25);
102
+ expect(result.slice(0, 25).every(({ sizeBytes }) => sizeBytes === 100)).toBe(true);
103
+ expect(result.slice(25).every(({ sizeBytes }) => sizeBytes === null)).toBe(true);
104
+ });
105
+ });
106
+ describe('BigqueryWarehouseClient.getProjectRecommendation', () => {
107
+ const createDataset = (projectId, datasetId, location) => new bigquery_1.Dataset(new bigquery_1.BigQuery({ projectId }), datasetId, {
108
+ location,
109
+ });
110
+ const createClient = (projectId, locations, queryImplementation) => {
111
+ const datasets = locations.map((location, index) => createDataset(projectId, `dataset_${index}`, location));
112
+ const getDatasets = vi.fn();
113
+ vi.mocked(getDatasets).mockResolvedValue([datasets]);
114
+ const query = vi.fn();
115
+ vi.mocked(query).mockImplementation(queryImplementation);
116
+ return { getDatasets, query };
117
+ };
118
+ it('ranks projects by their largest dataset across distinct regions', async () => {
119
+ const firstClient = createClient('first-project', ['EU', 'US', 'EU'], async ({ location }) => location === 'US'
120
+ ? [[{ table_schema: 'largest', total_bytes: 1200 }]]
121
+ : [[{ table_schema: 'smaller', total_bytes: 500 }]]);
122
+ const secondClient = createClient('second-project', ['EU'], async () => [
123
+ [
124
+ { table_schema: 'small', total_bytes: 100 },
125
+ { table_schema: 'large', total_bytes: 900 },
126
+ ],
127
+ ]);
128
+ const clients = new Map([
129
+ ['first-project', firstClient],
130
+ ['second-project', secondClient],
131
+ ]);
132
+ const clientFactory = vi.fn((projectId) => {
133
+ const client = clients.get(projectId);
134
+ if (!client)
135
+ throw new Error('Unexpected project');
136
+ return client;
137
+ });
138
+ const projects = [
139
+ { projectId: 'first-project', friendlyName: null },
140
+ { projectId: 'second-project', friendlyName: null },
141
+ ];
142
+ await expect(BigqueryWarehouseClient_1.BigqueryWarehouseClient.getProjectRecommendation(projects, 'refresh-token', clientFactory)).resolves.toEqual({ projectId: 'first-project' });
143
+ expect(vi.mocked(firstClient.getDatasets)).toHaveBeenCalledWith({
144
+ autoPaginate: false,
145
+ maxResults: 1000,
146
+ });
147
+ expect(vi.mocked(firstClient.query)).toHaveBeenCalledTimes(2);
148
+ expect(vi.mocked(firstClient.query)).toHaveBeenCalledWith({
149
+ query: 'SELECT table_schema, SUM(total_logical_bytes) AS total_bytes FROM `first-project`.`region-us`.INFORMATION_SCHEMA.TABLE_STORAGE GROUP BY table_schema',
150
+ location: 'US',
151
+ });
152
+ });
153
+ it('skips a project when any of its region queries errors', async () => {
154
+ const failingClient = createClient('failing-project', ['EU', 'US'], async ({ location }) => {
155
+ if (location === 'US')
156
+ throw new Error('permission denied');
157
+ return [[{ table_schema: 'largest', total_bytes: 5000 }]];
158
+ });
159
+ const healthyClient = createClient('healthy-project', ['EU'], async () => [[{ table_schema: 'available', total_bytes: 100 }]]);
160
+ const clients = new Map([
161
+ ['failing-project', failingClient],
162
+ ['healthy-project', healthyClient],
163
+ ]);
164
+ const clientFactory = vi.fn((projectId) => {
165
+ const client = clients.get(projectId);
166
+ if (!client)
167
+ throw new Error('Unexpected project');
168
+ return client;
169
+ });
170
+ await expect(BigqueryWarehouseClient_1.BigqueryWarehouseClient.getProjectRecommendation([
171
+ { projectId: 'failing-project', friendlyName: null },
172
+ { projectId: 'healthy-project', friendlyName: null },
173
+ ], 'refresh-token', clientFactory)).resolves.toEqual({ projectId: 'healthy-project' });
174
+ });
175
+ it('returns no recommendation when every project score is unavailable', async () => {
176
+ const client = createClient('unavailable-project', ['EU'], async () => {
177
+ throw new Error('permission denied');
178
+ });
179
+ const clientFactory = vi.fn();
180
+ vi.mocked(clientFactory).mockReturnValue(client);
181
+ await expect(BigqueryWarehouseClient_1.BigqueryWarehouseClient.getProjectRecommendation([{ projectId: 'unavailable-project', friendlyName: null }], 'refresh-token', clientFactory)).resolves.toEqual({ projectId: null });
182
+ });
183
+ it('caps ranking at eight projects and three regions per project', async () => {
184
+ const projects = Array.from({ length: 10 }, (_, index) => ({
185
+ projectId: `project-${index}`,
186
+ friendlyName: null,
187
+ }));
188
+ const clients = new Map(projects.map(({ projectId }, index) => [
189
+ projectId,
190
+ createClient(projectId, ['EU', 'US', 'asia-east1', 'us-west1'], async () => [
191
+ [{ table_schema: 'dataset', total_bytes: index }],
192
+ ]),
193
+ ]));
194
+ const clientFactory = vi.fn((projectId) => {
195
+ const client = clients.get(projectId);
196
+ if (!client)
197
+ throw new Error('Unexpected project');
198
+ return client;
199
+ });
200
+ await expect(BigqueryWarehouseClient_1.BigqueryWarehouseClient.getProjectRecommendation(projects, 'refresh-token', clientFactory)).resolves.toEqual({ projectId: 'project-7' });
201
+ expect(vi.mocked(clientFactory)).toHaveBeenCalledTimes(8);
202
+ projects.forEach(({ projectId }, index) => {
203
+ const client = clients.get(projectId);
204
+ if (!client)
205
+ throw new Error('Unexpected project');
206
+ expect(vi.mocked(client.getDatasets)).toHaveBeenCalledTimes(index < 8 ? 1 : 0);
207
+ expect(vi.mocked(client.query)).toHaveBeenCalledTimes(index < 8 ? 3 : 0);
208
+ });
209
+ });
210
+ });
45
211
  describe('BigqueryWarehouseClient.sanitizeLabelsWithValues', () => {
46
212
  let warnSpy;
47
213
  beforeEach(() => {
@@ -49,6 +49,7 @@ export type SnowflakeSessionInventory = {
49
49
  name: string;
50
50
  comment: string | null;
51
51
  kind: string | null;
52
+ sizeBytes: number | null;
52
53
  }[];
53
54
  warehouses: {
54
55
  name: string;
@@ -1 +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"}
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;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"}
@@ -96,6 +96,7 @@ const EXTERNAL_BROWSER_AUTHENTICATOR = 'EXTERNALBROWSER';
96
96
  const OAUTH_AUTHORIZATION_CODE_AUTHENTICATOR = 'OAUTH_AUTHORIZATION_CODE';
97
97
  const SESSION_DISCOVERY_LIMIT = 100;
98
98
  const SESSION_SCHEMA_DISCOVERY_LIMIT = 1000;
99
+ const SESSION_DATABASE_SIZE_LIMIT = 25;
99
100
  // Minimum allowed by Snowflake (range 16-160, default 160)
100
101
  const SNOWFLAKE_RESULT_CHUNK_SIZE_MB = 16;
101
102
  const getAccessTokenExpiresAt = (token) => {
@@ -202,7 +203,7 @@ const listSnowflakeSchemas = (rows) => rows.flatMap((rawRow) => {
202
203
  }
203
204
  return [{ databaseName, name }];
204
205
  });
205
- const listSnowflakeDatabases = (rows) => rows
206
+ const listSnowflakeDatabases = (rows, databaseSizes) => rows
206
207
  .flatMap((rawRow) => {
207
208
  const row = rawRow;
208
209
  const name = getSnowflakeRowValue(row, 'name');
@@ -216,6 +217,7 @@ const listSnowflakeDatabases = (rows) => rows
216
217
  ? null
217
218
  : getSnowflakeRowValue(row, 'comment'),
218
219
  kind: getSnowflakeRowValue(row, 'kind'),
220
+ sizeBytes: databaseSizes.get(name) ?? null,
219
221
  },
220
222
  ];
221
223
  })
@@ -237,6 +239,44 @@ const listSnowflakeWarehouses = (rows) => rows
237
239
  ];
238
240
  })
239
241
  .slice(0, SESSION_DISCOVERY_LIMIT);
242
+ const WAREHOUSE_SIZE_ORDER = [
243
+ 'X-Small',
244
+ 'Small',
245
+ 'Medium',
246
+ 'Large',
247
+ 'X-Large',
248
+ '2X-Large',
249
+ '3X-Large',
250
+ '4X-Large',
251
+ '5X-Large',
252
+ '6X-Large',
253
+ ];
254
+ const warehouseSizeRank = (size) => {
255
+ const index = WAREHOUSE_SIZE_ORDER.findIndex((candidate) => candidate.toLowerCase() === (size ?? '').trim().toLowerCase());
256
+ return index === -1 ? Number.MAX_SAFE_INTEGER : index;
257
+ };
258
+ const pickSizeDiscoveryWarehouse = (rows) => {
259
+ const warehouses = rows
260
+ .map((rawRow) => {
261
+ const row = rawRow;
262
+ return {
263
+ name: getSnowflakeRowValue(row, 'name'),
264
+ size: getSnowflakeRowValue(row, 'size'),
265
+ state: (getSnowflakeRowValue(row, 'state') ?? '')
266
+ .trim()
267
+ .toUpperCase(),
268
+ };
269
+ })
270
+ .filter((warehouse) => Boolean(warehouse.name));
271
+ const started = warehouses.filter((warehouse) => warehouse.state === 'STARTED');
272
+ const pool = started.length > 0 ? started : warehouses;
273
+ const best = pool.reduce((currentBest, candidate) => currentBest === null ||
274
+ warehouseSizeRank(candidate.size) <
275
+ warehouseSizeRank(currentBest.size)
276
+ ? candidate
277
+ : currentBest, null);
278
+ return best?.name ?? null;
279
+ };
240
280
  const snowflakeIdentifier = (value) => /^[A-Za-z_][A-Za-z0-9_$]*$/.test(value)
241
281
  ? value
242
282
  : `"${value.replace(/"/g, '""')}"`;
@@ -623,6 +663,38 @@ class SnowflakeWarehouseClient extends WarehouseBaseClient_1.default {
623
663
  catch {
624
664
  schemaRows = [];
625
665
  }
666
+ const databaseSizes = new Map();
667
+ const sizeCandidates = databasesResult.rows
668
+ .map((rawRow) => getSnowflakeRowValue(rawRow, 'name'))
669
+ .filter((name) => Boolean(name))
670
+ .slice(0, SESSION_DATABASE_SIZE_LIMIT);
671
+ // INFORMATION_SCHEMA aggregation needs an active warehouse; without a
672
+ // session default, activate one (prefer already-running, then smallest)
673
+ if (sizeCandidates.length > 0) {
674
+ const sessionWarehouseName = getSnowflakeRowValue(defaultsRow, 'warehouse');
675
+ const sizeWarehouseName = sessionWarehouseName ??
676
+ pickSizeDiscoveryWarehouse(warehousesResult.rows);
677
+ if (!sessionWarehouseName && sizeWarehouseName) {
678
+ try {
679
+ await this.executeStatements(connection, `USE WAREHOUSE ${(0, exports.snowflakeIdentifier)(sizeWarehouseName)}`);
680
+ }
681
+ catch {
682
+ /* size queries will fail and sizes stay null */
683
+ }
684
+ }
685
+ }
686
+ await sizeCandidates.reduce((previous, databaseName) => previous.then(async () => {
687
+ try {
688
+ 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'`);
689
+ const totalBytes = getSnowflakeRowNumber((sizeResult.rows[0] ?? {}), 'total_bytes');
690
+ if (totalBytes !== null && totalBytes > 0) {
691
+ databaseSizes.set(databaseName, totalBytes);
692
+ }
693
+ }
694
+ catch {
695
+ /* size stays null for this database */
696
+ }
697
+ }), Promise.resolve());
626
698
  const defaultRole = getSnowflakeRowValue(defaultsRow, 'role');
627
699
  const grantedRoles = [
628
700
  ...new Set([
@@ -649,7 +721,7 @@ class SnowflakeWarehouseClient extends WarehouseBaseClient_1.default {
649
721
  schema: getSnowflakeRowValue(defaultsRow, 'schema'),
650
722
  },
651
723
  inventory: {
652
- databases: listSnowflakeDatabases(databasesResult.rows),
724
+ databases: listSnowflakeDatabases(databasesResult.rows, databaseSizes),
653
725
  warehouses: listSnowflakeWarehouses(warehousesResult.rows),
654
726
  roles,
655
727
  schemas: listSnowflakeSchemas(schemaRows),
@@ -311,6 +311,17 @@ describe('SnowflakeWarehouseClient', () => {
311
311
  },
312
312
  ];
313
313
  }
314
+ else if (sqlText.startsWith('SELECT SUM(bytes)')) {
315
+ if (sqlText.includes('DATABASE_0.')) {
316
+ rows = [{ total_bytes: 1250 }];
317
+ }
318
+ else if (sqlText.includes('DATABASE_1.')) {
319
+ rows = [{ total_bytes: 500 }];
320
+ }
321
+ else {
322
+ rows = [{ total_bytes: null }];
323
+ }
324
+ }
314
325
  complete(undefined, { getColumns: () => SnowflakeWarehouseClient_mock_1.queryColumnsMock }, rows);
315
326
  });
316
327
  vi.mocked(snowflake_sdk_1.createConnection).mockImplementationOnce(() => interactiveConnectionMock(execute));
@@ -331,6 +342,7 @@ describe('SnowflakeWarehouseClient', () => {
331
342
  name: `DATABASE_${index}`,
332
343
  comment: index === 0 ? 'Primary analytics' : null,
333
344
  kind: index === 1 ? 'IMPORTED DATABASE' : 'STANDARD',
345
+ sizeBytes: [1250, 500][index] ?? null,
334
346
  })),
335
347
  warehouses: [
336
348
  {
@@ -354,10 +366,56 @@ describe('SnowflakeWarehouseClient', () => {
354
366
  expect(execute).toHaveBeenCalledWith(expect.objectContaining({
355
367
  sqlText: 'SHOW SCHEMAS IN ACCOUNT LIMIT 1000',
356
368
  }));
369
+ expect(execute).toHaveBeenCalledWith(expect.objectContaining({
370
+ sqlText: `SELECT SUM(bytes) AS "total_bytes" FROM DATABASE_0.INFORMATION_SCHEMA.TABLES WHERE table_type = 'BASE TABLE'`,
371
+ }));
372
+ expect(execute.mock.calls.filter(([statement]) => statement.sqlText.startsWith('SELECT SUM(bytes)'))).toHaveLength(25);
373
+ expect(execute.mock.calls.some(([statement]) => statement.sqlText.startsWith('USE WAREHOUSE'))).toBe(false);
357
374
  expect(execute).toHaveBeenCalledWith(expect.objectContaining({
358
375
  sqlText: 'SHOW GRANTS TO USER "user.name@example.com" LIMIT 100',
359
376
  }));
360
377
  });
378
+ it('returns null database sizes when size discovery fails', async () => {
379
+ const execute = vi.fn(({ sqlText, complete, }) => {
380
+ if (sqlText.startsWith('SELECT SUM(bytes)')) {
381
+ complete(new Error('Insufficient privileges'), { getColumns: () => SnowflakeWarehouseClient_mock_1.queryColumnsMock }, []);
382
+ return;
383
+ }
384
+ let rows = [];
385
+ if (sqlText.startsWith('SELECT CURRENT_USER')) {
386
+ rows = [{ USER: 'test-user' }];
387
+ }
388
+ else if (sqlText.startsWith('SHOW DATABASES')) {
389
+ rows = [{ name: 'ANALYTICS' }, { name: 'RAW' }];
390
+ }
391
+ else if (sqlText.startsWith('SHOW WAREHOUSES')) {
392
+ rows = [
393
+ { name: 'WH_BIG', size: 'Large', state: 'SUSPENDED' },
394
+ {
395
+ name: 'WH_SMALL',
396
+ size: 'X-Small',
397
+ state: 'STARTED',
398
+ },
399
+ ];
400
+ }
401
+ complete(undefined, { getColumns: () => SnowflakeWarehouseClient_mock_1.queryColumnsMock }, rows);
402
+ });
403
+ vi.mocked(snowflake_sdk_1.createConnection).mockImplementationOnce(() => interactiveConnectionMock(execute));
404
+ const warehouse = new SnowflakeWarehouseClient_1.SnowflakeWarehouseClient(SnowflakeWarehouseClient_mock_1.credentials);
405
+ const discovery = await warehouse.getSessionDiscovery('test-user');
406
+ expect(discovery.inventory.databases).toEqual([
407
+ {
408
+ name: 'ANALYTICS',
409
+ comment: null,
410
+ kind: null,
411
+ sizeBytes: null,
412
+ },
413
+ { name: 'RAW', comment: null, kind: null, sizeBytes: null },
414
+ ]);
415
+ expect(execute).toHaveBeenCalledWith(expect.objectContaining({
416
+ sqlText: 'USE WAREHOUSE WH_SMALL',
417
+ }));
418
+ });
361
419
  it('opens, checks, and closes a diagnostic connection', async () => {
362
420
  const execute = vi.fn(({ sqlText, complete, }) => complete(undefined, { getColumns: () => SnowflakeWarehouseClient_mock_1.queryColumnsMock }, []));
363
421
  const connection = interactiveConnectionMock(execute);
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@lightdash/warehouses",
3
- "version": "0.3413.1",
3
+ "version": "0.3415.0",
4
4
  "description": "Warehouse connectors for Lightdash",
5
5
  "license": "MIT",
6
6
  "repository": {
@@ -30,7 +30,7 @@
30
30
  "snowflake-sdk": "2.3.4",
31
31
  "ssh2": "1.14.0",
32
32
  "trino-client": "0.2.9",
33
- "@lightdash/common": "0.3413.1"
33
+ "@lightdash/common": "0.3415.0"
34
34
  },
35
35
  "devDependencies": {
36
36
  "@types/big.js": "6.2.2",