@malloy-publisher/server 0.0.165 → 0.0.168

package/src/service/db_utils.ts

@@ -7,9 +7,8 @@ import {
    CloudStorageCredentials,
    gcsConnectionToCredentials,
    getCloudTablesWithColumns,
-   listAllDataFilesInBucket,
-   listCloudBuckets,
-   listFilesInCloudDirectory,
+   listCloudDirectorySchemas,
+   listDataFilesInDirectory,
    parseCloudUri,
    s3ConnectionToCredentials,
 } from "./gcs_s3_utils";
@@ -340,23 +339,10 @@ export async function getSchemasForConnection(
                   : s3ConnectionToCredentials(attachedDb.s3Connection!);
 
             try {
-               const buckets = await listCloudBuckets(credentials);
-               const scheme = dbType === "gcs" ? "gs" : "s3";
-
-               logger.info(
-                  `Listed ${buckets.length} ${dbType.toUpperCase()} buckets for attached database ${attachedDb.name}`,
-               );
-
-               // Just return bucket URIs as schemas - fast!
-               // Files/directories will be listed when user selects a bucket
-               return buckets.map((bucket) => ({
-                  name: `${scheme}://${bucket.name}`,
-                  isHidden: false,
-                  isDefault: false,
-               }));
+               return await listCloudDirectorySchemas(credentials);
             } catch (cloudError) {
                logger.warn(
-                  `Failed to list ${dbType.toUpperCase()} buckets for ${attachedDb.name}`,
+                  `Failed to list ${dbType.toUpperCase()} directory schemas for ${attachedDb.name}`,
                   { error: cloudError },
                );
                return [];
@@ -411,6 +397,38 @@ export async function getSchemasForConnection(
             `Failed to get schemas for MotherDuck connection ${connection.name}: ${(error as Error).message}`,
          );
       }
+   } else if (connection.type === "ducklake") {
+      try {
+         // Filter by catalog_name to only get schemas from the attached DuckLake catalog
+         // The catalog is attached with the connection name (see attachDuckLake in connection.ts)
+         const catalogName = connection.name;
+         const result = await malloyConnection.runSQL(
+            `SELECT schema_name FROM information_schema.schemata WHERE catalog_name = '${catalogName}' ORDER BY schema_name`,
+            { rowLimit: 1000 },
+         );
+         const rows = standardizeRunSQLResult(result);
+
+         return rows.map((row: unknown) => {
+            const typedRow = row as Record<string, unknown>;
+            const schemaName = typedRow.schema_name as string;
+
+            const shouldShow = schemaName === "main" || schemaName === "public";
+
+            return {
+               name: schemaName,
+               isHidden: !shouldShow,
+               isDefault: false,
+            };
+         });
+      } catch (error) {
+         logger.error(
+            `Error getting schemas for DuckLake connection ${connection.name}`,
+            { error },
+         );
+         throw new Error(
+            `Failed to get schemas for DuckLake connection ${connection.name}: ${(error as Error).message}`,
+         );
+      }
    } else {
       throw new Error(`Unsupported connection type: ${connection.type}`);
    }
@@ -444,19 +462,11 @@ export async function getTablesForSchema(
          );
       }
 
-      let fileKeys: string[];
-      if (directoryPath) {
-         const fileNames = await listFilesInCloudDirectory(
-            credentials,
-            bucketName,
-            directoryPath,
-         );
-         fileKeys = fileNames.map((fileName) => `${directoryPath}/${fileName}`);
-      } else {
-         fileKeys = await listAllDataFilesInBucket(credentials, bucketName);
-      }
-
-      console.log("File keys:", fileKeys);
+      const fileKeys = await listDataFilesInDirectory(
+         credentials,
+         bucketName,
+         directoryPath,
+      );
 
       return await getCloudTablesWithColumns(
          malloyConnection,
@@ -464,8 +474,13 @@ export async function getTablesForSchema(
          bucketName,
          fileKeys,
       );
+   } else if (connection.type === "ducklake") {
+      if (schemaName.split(".").length == 2) {
+         schemaName = `${connection.name}.${schemaName}`;
+      } else if (schemaName.split(".").length === 1) {
+         schemaName = `${connection.name}.${schemaName}`;
+      }
    }
-
    const tableNames = await listTablesForSchema(
       connection,
       schemaName,
@@ -476,7 +491,6 @@ export async function getTablesForSchema(
    const tableSourcePromises = tableNames.map(async (tableName) => {
       try {
          let tablePath: string;
-
          if (connection.type === "trino") {
             if (connection.trinoConnection?.catalog) {
                tablePath = `${connection.trinoConnection?.catalog}.${schemaName}.${tableName}`;
@@ -484,6 +498,10 @@ export async function getTablesForSchema(
                // Catalog name is included in the schema name
                tablePath = `${schemaName}.${tableName}`;
             }
+         } else if (connection.type === "ducklake") {
+            // For ducklake, schemaName already includes connection name prefix from above
+            // So tablePath should be schemaName.tableName (which is connectionName.schemaName.tableName)
+            tablePath = `${schemaName}.${tableName}`;
          } else {
             tablePath = `${schemaName}.${tableName}`;
          }
@@ -497,14 +515,13 @@ export async function getTablesForSchema(
             tableName,
             tablePath,
          );
-
          return {
             resource: tablePath,
             columns: tableSource.columns,
          };
       } catch (error) {
          logger.warn(`Failed to get schema for table ${tableName}`, {
-            error,
+            error: extractErrorDataFromError(error),
             schemaName,
             tableName,
          });
@@ -575,7 +592,7 @@ export async function getConnectionTableSource(
             type: field.type,
          };
       });
-      logger.info(`Successfully fetched schema for ${tablePath}`, {
+      logger.debug(`Successfully fetched schema for ${tablePath}`, {
          fieldCount: fields.length,
       });
       return {
@@ -748,15 +765,15 @@ export async function listTablesForSchema(
          }
 
          try {
-            if (directoryPath) {
-               return await listFilesInCloudDirectory(
-                  credentials,
-                  bucketName,
-                  directoryPath,
-               );
-            } else {
-               return await listAllDataFilesInBucket(credentials, bucketName);
-            }
+            const fileKeys = await listDataFilesInDirectory(
+               credentials,
+               bucketName,
+               directoryPath,
+            );
+            return fileKeys.map((key) => {
+               const lastSlash = key.lastIndexOf("/");
+               return lastSlash > 0 ? key.substring(lastSlash + 1) : key;
+            });
          } catch (error) {
             logger.error(
                `Error listing ${cloudType.toUpperCase()} objects in ${schemaName}`,
@@ -817,7 +834,49 @@ export async function listTablesForSchema(
             `Failed to get tables for MotherDuck schema ${schemaName} in connection ${connection.name}: ${(error as Error).message}`,
          );
       }
+   } else if (connection.type === "ducklake") {
+      const catalogName = schemaName.split(".")[0];
+      const actualSchemaName = schemaName.split(".")[1];
+      console.error("catalogName", catalogName);
+      console.error("actualSchemaName", actualSchemaName);
+      try {
+         const result = await malloyConnection.runSQL(
+            `SELECT table_name FROM information_schema.tables WHERE table_schema = '${actualSchemaName}' AND table_catalog = '${catalogName}' ORDER BY table_name`,
+            { rowLimit: 1000 },
+         );
+         const rows = standardizeRunSQLResult(result);
+         return rows.map((row: unknown) => {
+            const typedRow = row as Record<string, unknown>;
+            return typedRow.table_name as string;
+         });
+      } catch (error) {
+         logger.error(
+            `Error getting tables for DuckLake schema ${schemaName} in connection ${connection.name}`,
+            { error },
+         );
+         throw new Error(
+            `Failed to get tables for DuckLake schema ${schemaName} in connection ${connection.name}: ${(error as Error).message}`,
+         );
+      }
    } else {
       throw new Error(`Unsupported connection type: ${connection.type}`);
    }
 }
+
+export function extractErrorDataFromError(error: unknown): {
+   error: string;
+   stack?: string;
+   task?: unknown;
+} {
+   const errorMessage = error instanceof Error ? error.message : String(error);
+   const errorData: { error: string; stack?: string; task?: unknown } = {
+      error: errorMessage,
+   };
+   if (error instanceof Error && logger.level === "debug") {
+      errorData.stack = error.stack;
+   }
+   if (error && typeof error === "object" && "task" in error) {
+      errorData.task = (error as { task?: unknown }).task;
+   }
+   return errorData;
+}

package/src/service/gcs_s3_utils.ts

@@ -8,7 +8,6 @@ import { components } from "../api";
 import { logger } from "../logger";
 
 type ApiTable = components["schemas"]["Table"];
-
 type CloudStorageType = "gcs" | "s3";
 
 export interface CloudStorageCredentials {
@@ -29,7 +28,6 @@ interface CloudStorageObject {
    key: string;
    size?: number;
    lastModified?: Date;
-   isFolder: boolean;
 }
 
 export function gcsConnectionToCredentials(gcsConnection: {
@@ -92,7 +90,7 @@ function createCloudStorageClient(
    return client;
 }
 
-export async function listCloudBuckets(
+async function listCloudBuckets(
    credentials: CloudStorageCredentials,
 ): Promise<CloudStorageBucket[]> {
    const client = createCloudStorageClient(credentials);
@@ -143,7 +141,6 @@ async function listAllCloudFiles(
                   key: content.Key,
                   size: content.Size,
                   lastModified: content.LastModified,
-                  isFolder: false,
                });
             }
          }
@@ -182,6 +179,15 @@ function isDataFile(key: string): boolean {
    );
 }
 
+function buildCloudUri(
+   type: CloudStorageType,
+   bucket: string,
+   key: string,
+): string {
+   const scheme = type === "gcs" ? "gs" : "s3";
+   return `${scheme}://${bucket}/${key}`;
+}
+
 function getFileType(key: string): string {
    const lowerKey = key.toLowerCase();
    if (lowerKey.endsWith(".csv")) return "csv";
@@ -192,23 +198,14 @@ function getFileType(key: string): string {
    return "unknown";
 }
 
-function buildCloudUri(
-   type: CloudStorageType,
-   bucket: string,
-   key: string,
-): string {
-   const scheme = type === "gcs" ? "gs" : "s3";
-   return `${scheme}://${bucket}/${key}`;
-}
-
 function standardizeRunSQLResult(result: unknown): unknown[] {
    return Array.isArray(result)
       ? result
       : (result as { rows?: unknown[] }).rows || [];
 }
 
-// Batch size for parallel schema fetching to avoid overwhelming the connection
 const SCHEMA_FETCH_BATCH_SIZE = 10;
+const BUCKET_SCAN_BATCH_SIZE = 3;
 
 async function getTableSchema(
    malloyConnection: Connection,
@@ -268,11 +265,9 @@ export async function getCloudTablesWithColumns(
 ): Promise<ApiTable[]> {
    const allTables: ApiTable[] = [];
 
-   // Process in batches to avoid overwhelming the connection
    for (let i = 0; i < fileKeys.length; i += SCHEMA_FETCH_BATCH_SIZE) {
       const batch = fileKeys.slice(i, i + SCHEMA_FETCH_BATCH_SIZE);
 
-      // Process batch in parallel
       const batchResults = await Promise.all(
          batch.map((fileKey) =>
             getTableSchema(malloyConnection, credentials, bucketName, fileKey),
@@ -315,38 +310,118 @@ export function parseCloudUri(uri: string): {
    return null;
 }
 
-export async function listFilesInCloudDirectory(
+export async function listDataFilesInDirectory(
    credentials: CloudStorageCredentials,
    bucketName: string,
    directoryPath: string,
 ): Promise<string[]> {
-   const files = await listAllCloudFiles(credentials, bucketName);
-
-   const filesInDirectory = files
-      .filter((obj) => {
-         if (!isDataFile(obj.key)) return false;
-
-         const lastSlashIndex = obj.key.lastIndexOf("/");
-         const fileDir =
-            lastSlashIndex > 0 ? obj.key.substring(0, lastSlashIndex) : "";
-
-         return fileDir === directoryPath;
-      })
-      .map((obj) => {
-         const lastSlashIndex = obj.key.lastIndexOf("/");
-         return lastSlashIndex > 0
-            ? obj.key.substring(lastSlashIndex + 1)
-            : obj.key;
-      });
+   const prefix = directoryPath ? `${directoryPath}/` : "";
+   const client = createCloudStorageClient(credentials);
+   const storageType = credentials.type.toUpperCase();
+   const dataFiles: string[] = [];
+
+   try {
+      let continuationToken: string | undefined;
+
+      do {
+         const response = await client.send(
+            new ListObjectsV2Command({
+               Bucket: bucketName,
+               Prefix: prefix,
+               Delimiter: "/",
+               ContinuationToken: continuationToken,
+            }),
+         );
+
+         for (const content of response.Contents || []) {
+            if (content.Key && isDataFile(content.Key)) {
+               dataFiles.push(content.Key);
+            }
+         }
 
-   return filesInDirectory;
+         continuationToken = response.IsTruncated
+            ? response.NextContinuationToken
+            : undefined;
+      } while (continuationToken);
+
+      logger.info(
+         `Listed ${dataFiles.length} data files in ${storageType} ${bucketName}/${directoryPath}`,
+      );
+      return dataFiles;
+   } catch (error) {
+      logger.error(
+         `Failed to list files in ${storageType} ${bucketName}/${directoryPath}`,
+         { error },
+      );
+      throw new Error(
+         `Failed to list files in ${storageType} ${bucketName}/${directoryPath}: ${error instanceof Error ? error.message : String(error)}`,
+      );
+   }
 }
 
-// List all data files in a bucket with their full relative paths
-export async function listAllDataFilesInBucket(
+/**
+ * Scans an entire bucket and returns unique directory paths that contain data files.
+ * Uses flat listing for efficiency — O(total_files / 1000) API calls.
+ */
+async function listDirectorySchemas(
    credentials: CloudStorageCredentials,
    bucketName: string,
 ): Promise<string[]> {
-   const files = await listAllCloudFiles(credentials, bucketName);
-   return files.filter((obj) => isDataFile(obj.key)).map((obj) => obj.key);
+   const allFiles = await listAllCloudFiles(credentials, bucketName);
+   const directories = new Set<string>();
+
+   for (const file of allFiles) {
+      if (!isDataFile(file.key)) continue;
+
+      const lastSlashIndex = file.key.lastIndexOf("/");
+      const dir =
+         lastSlashIndex > 0 ? file.key.substring(0, lastSlashIndex) : "";
+      directories.add(dir);
+   }
+
+   const scheme = credentials.type === "gcs" ? "gs" : "s3";
+   const sortedDirs = Array.from(directories).sort();
+
+   logger.info(
+      `Found ${sortedDirs.length} directories with data files in ${credentials.type.toUpperCase()} bucket ${bucketName}`,
+   );
+
+   return sortedDirs.map((dir) =>
+      dir ? `${scheme}://${bucketName}/${dir}` : `${scheme}://${bucketName}`,
+   );
+}
+
+export async function listCloudDirectorySchemas(
+   credentials: CloudStorageCredentials,
+): Promise<{ name: string; isHidden: boolean; isDefault: boolean }[]> {
+   const storageType = credentials.type.toUpperCase();
+   const buckets = await listCloudBuckets(credentials);
+
+   logger.info(
+      `Listed ${buckets.length} ${storageType} buckets, scanning for directories...`,
+   );
+
+   const allDirArrays: string[][] = [];
+
+   for (let i = 0; i < buckets.length; i += BUCKET_SCAN_BATCH_SIZE) {
+      const batch = buckets.slice(i, i + BUCKET_SCAN_BATCH_SIZE);
+      const batchResults = await Promise.all(
+         batch.map((bucket) =>
+            listDirectorySchemas(credentials, bucket.name).catch((err) => {
+               logger.warn(
+                  `Failed to scan ${storageType} bucket ${bucket.name}`,
+                  { error: err },
+               );
+               return [] as string[];
+            }),
+         ),
+      );
+      allDirArrays.push(...batchResults);
+   }
+
+   return allDirArrays.flat().map((dirUri) => ({
+      name: dirUri,
+      isHidden: false,
+      isDefault: false,
+   }));
 }

package/src/service/model.ts

@@ -211,13 +211,13 @@ export class Model {
       } catch (error) {
          let computedError = error;
          if (error instanceof Error && error.stack) {
-            console.error("Error stack", error.stack);
+            logger.error("Error stack", error.stack);
          }
 
          if (error instanceof MalloyError) {
             const problems = error.problems;
             for (const problem of problems) {
-               console.error("Problem", problem);
+               logger.error("Problem", problem);
             }
             computedError = new ModelCompilationError(error);
          }
@@ -444,7 +444,7 @@ export class Model {
       const notebookCells: ApiNotebookCell[] = (
          this.runnableNotebookCells as RunnableNotebookCell[]
       ).map((cell) => {
-         console.log("cell.queryInfo", cell.queryInfo);
+         logger.debug("cell.queryInfo", cell.queryInfo);
          return {
             type: cell.type,
             text: cell.text,
@@ -550,9 +550,9 @@ export class Model {
                   text: cell.text,
                };
             } else {
-               console.log("Error message: ", errorMessage);
+               logger.error("Error message: ", errorMessage);
             }
-            console.log("Cell content: ", cellIndex, cell.type, cell.text);
+            logger.debug("Cell content: ", cellIndex, cell.type, cell.text);
             throw new BadRequestError(`Cell execution failed: ${errorMessage}`);
          }
       }