snowflake filesystem (#39)

Author: sfc-gh-hyu

plugins/nf-snowflake/build.gradle

@@ -87,4 +87,3 @@ dependencies {
 test {
     useJUnitPlatform()
 }
-

plugins/nf-snowflake/src/main/nextflow/snowflake/SnowflakeCacheFactory.groovy

@@ -10,6 +10,7 @@ import nextflow.exception.AbortOperationException
 import nextflow.plugin.Priority
 import nextflow.cache.CacheDB
 import nextflow.cache.CacheFactory
+import nextflow.cache.DefaultCacheStore
 
 /**
  * Implements the cloud cache factory
@@ -27,11 +28,17 @@ class SnowflakeCacheFactory extends CacheFactory {
         if( !uniqueId ) throw new AbortOperationException("Missing cache `uuid`")
         if( !runName ) throw new AbortOperationException("Missing cache `runName`")
 
-        // Defer validation of SNOWFLAKE_CACHE_PATH until the cache is actually used
+        // Use SNOWFLAKE_CACHE_PATH if set, otherwise fall back to DefaultCacheStore
         final String cachePathEnv = System.getenv("SNOWFLAKE_CACHE_PATH")
-        final Path basePath = cachePathEnv ? Paths.get(cachePathEnv) : null
 
-        final store = new SnowflakeCacheStore(uniqueId, runName, basePath)
-        return new CacheDB(store)
+        if( cachePathEnv ) {
+            final Path basePath = Paths.get(cachePathEnv)
+            final store = new SnowflakeCacheStore(uniqueId, runName, basePath)
+            return new CacheDB(store)
+        }
+        else {
+            final store = new DefaultCacheStore(uniqueId, runName, home)
+            return new CacheDB(store)
+        }
     }
 }

plugins/nf-snowflake/src/main/nextflow/snowflake/SnowflakeExecutor.groovy

@@ -85,24 +85,31 @@ class SnowflakeExecutor extends Executor implements ExtensionPoint {
     }
 
     /**
-     * Copy local bin directory to remote mounted workdir
+     * Copy local bin directory to Snowflake stage
      */
     protected void uploadBinDir() {
         if( session.binDir && !session.binDir.empty() && !session.disableRemoteBinDir ) {
             // Use session run name for directory isolation
             final String runId = session.runName
-            final Path targetDir = Paths.get("/mnt/workdir", runId, "bin")
-            
-            // Create target directory
+            final Path workDirPath = session.workDir
+
+            log.debug("Uploading bin directory to Snowflake stage: ${workDirPath.toUriString()}")
+
+            // Upload to snowflake://stage/STAGE_NAME/<runId>/bin
+            final Path targetDir = workDirPath.resolve(runId).resolve("bin")
+
+            // Create target directory (implicit in Snowflake)
             Files.createDirectories(targetDir)
-            
-            // Copy all files from bin directory using standard file I/O
+
+            // Copy all files from bin directory to stage
             Files.list(session.binDir).forEach { Path source ->
                 final Path target = targetDir.resolve(source.getFileName())
                 Files.copy(source, target, StandardCopyOption.REPLACE_EXISTING)
+                log.debug("Uploaded bin file: ${source.getFileName()} -> ${target.toUriString()}")
             }
-            
+
             remoteBinDir = targetDir
+            log.debug("Bin directory uploaded to: ${targetDir.toUriString()}")
         }
     }
 
@@ -113,9 +120,40 @@ class SnowflakeExecutor extends Executor implements ExtensionPoint {
     protected void register() {
         super.register()
         snowflakeConfig = session.config.navigate("snowflake") as Map
+
+        // Validate that workDir uses snowflake:// scheme
+        validateWorkDir()
+
         uploadBinDir()
     }
 
+    /**
+     * Validate that the work directory uses the snowflake:// scheme
+     * @throws IllegalArgumentException if workDir doesn't use snowflake:// scheme
+     */
+    private void validateWorkDir() {
+        final Path workDirPath = session.workDir
+        if (!workDirPath) {
+            throw new IllegalArgumentException(
+                "Work directory is not configured. " +
+                "For Snowflake executor, you must specify a work directory using the snowflake:// scheme, " +
+                "e.g., workDir = 'snowflake://stage/MY_STAGE/work'"
+            )
+        }
+
+        final String workDirStr = workDirPath.toUriString()
+        if (!SnowflakeUri.isSnowflakeStageUri(workDirStr)) {
+            throw new IllegalArgumentException(
+                "Invalid work directory for Snowflake executor: ${workDirStr}\n" +
+                "The work directory must use the snowflake:// scheme pointing to a Snowflake internal stage.\n" +
+                "Example: workDir = 'snowflake://stage/MY_STAGE/work'\n" +
+                "Current workDir: ${workDirStr}"
+            )
+        }
+
+        log.debug("Using Snowflake stage as work directory: ${workDirStr}")
+    }
+
     @Override
     boolean isContainerNative() {
         return true

plugins/nf-snowflake/src/main/nextflow/snowflake/SnowflakeFileCopyStrategy.groovy

@@ -1,15 +1,21 @@
 package nextflow.snowflake
 
 import groovy.transform.CompileStatic
+import groovy.util.logging.Slf4j
 import nextflow.processor.TaskBean
 import nextflow.executor.SimpleFileCopyStrategy
 import java.nio.file.Path
 
 /**
  * File copy strategy for Snowflake executor with remote bin directory support
  * 
+ * Note: Path translation is handled by SnowflakeWrapperBuilder.create() which
+ * pre-translates all paths in the TaskBean before this strategy is created.
+ * This class only needs to handle bin directory setup.
+ * 
  * @author Hongye Yu
  */
+@Slf4j
 @CompileStatic
 class SnowflakeFileCopyStrategy extends SimpleFileCopyStrategy {
 
@@ -22,21 +28,32 @@ class SnowflakeFileCopyStrategy extends SimpleFileCopyStrategy {
 
     /**
      * Override to prepend remote bin directory setup script
+     * 
+     * The remoteBinDir is a snowflake:// path, so we need to translate it
+     * to the container mount path for the script.
      */
     @Override
     String getEnvScript(Map environment, boolean container) {
-        if( remoteBinDir == null )
+        if( remoteBinDir == null ) {
+            log.debug("No remote bin directory configured")
             return super.getEnvScript(environment, container)
+        }
+        
+        // Translate bin directory path from snowflake:// to /mnt/stage/...
+        String translatedBinDir = SnowflakeUri.translateToMount(remoteBinDir).toUriString()
+        log.debug("Remote bin directory: ${remoteBinDir} -> ${translatedBinDir}")
 
         final script = """\
             # Copy and setup remote bin directory
             NXF_BIN=\$(mktemp -d)
-            cp -r ${remoteBinDir}/* \$NXF_BIN/
-            chmod +x \$NXF_BIN/*
+            cp -r ${translatedBinDir}/* \$NXF_BIN/ 2>&1 || echo "[DEBUG] Failed to copy bin files" >&2
+            chmod +x \$NXF_BIN/* 2>&1
             export PATH=\$NXF_BIN:\$PATH
             """.stripIndent()
 
-        return script + super.getEnvScript(environment, container)
+        // Get parent env script and handle null return value
+        def parentScript = super.getEnvScript(environment, container)
+        return script + (parentScript ?: '')
     }
 }
 

plugins/nf-snowflake/src/main/nextflow/snowflake/SnowflakePathFactory.groovy

@@ -0,0 +1,94 @@
+/*
+ * Copyright 2021, Seqera Labs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package nextflow.snowflake
+
+import java.nio.file.Path
+import java.nio.file.spi.FileSystemProvider
+
+import groovy.transform.CompileStatic
+import groovy.util.logging.Slf4j
+import nextflow.file.FileSystemPathFactory
+import nextflow.snowflake.nio.SnowflakePath
+import nextflow.snowflake.nio.SnowflakeFileSystemProvider
+
+/**
+ * Implements FileSystemPathFactory for Snowflake stages
+ *
+ * @author Hongye Yu
+ */
+@Slf4j
+@CompileStatic
+class SnowflakePathFactory extends FileSystemPathFactory {
+
+    @Override
+    protected Path parseUri(String uri) {
+        // Use centralized URI checking
+        if (!SnowflakeUri.isSnowflakeStageUri(uri)) {
+            return null
+        }
+
+        try {
+            URI parsedUri = new URI(uri)
+
+            // Find the Snowflake FileSystemProvider
+            SnowflakeFileSystemProvider provider = FileSystemProvider.installedProviders()
+                .find { it instanceof SnowflakeFileSystemProvider } as SnowflakeFileSystemProvider
+
+            if (!provider) {
+                log.warn("Snowflake FileSystemProvider not found")
+                return null
+            }
+
+            return provider.getPath(parsedUri)
+        }
+        catch (Exception e) {
+            log.warn("Failed to parse Snowflake URI: ${uri}", e)
+            return null
+        }
+    }
+
+    @Override
+    protected String toUriString(Path path) {
+        if (!(path instanceof SnowflakePath)) {
+            return null
+        }
+
+        SnowflakePath snowPath = (SnowflakePath) path
+        // Handle relative paths - return null to let default handler manage them
+        if (!snowPath.isAbsolute()) {
+            return null
+        }
+        return snowPath.toUri().toString()
+    }
+
+    @Override
+    protected String getBashLib(Path path) {
+        // No special bash library needed for Snowflake
+        return null
+    }
+
+    @Override
+    protected String getUploadCmd(String source, Path target) {
+        if (!(target instanceof SnowflakePath)) {
+            return null
+        }
+
+        SnowflakePath snowPath = (SnowflakePath) target
+        String stageRef = snowPath.toStageReference()
+        return "PUT file://${source} @${stageRef}"
+    }
+}

plugins/nf-snowflake/src/main/nextflow/snowflake/SnowflakePlugin.groovy

@@ -17,14 +17,19 @@
 package nextflow.snowflake
 
 import groovy.transform.CompileStatic
+import groovy.util.logging.Slf4j
+import nextflow.file.FileHelper
 import nextflow.plugin.BasePlugin
+import nextflow.snowflake.nio.SnowflakeFileSystemProvider
 import org.pf4j.PluginWrapper
 
 /**
- * Implements the Hello plugins entry point
+ * Nextflow plugin for Snowflake extensions
  *
  * @author Paolo Di Tommaso <paolo.ditommaso@gmail.com>
+ * @author Hongye Yu
  */
+@Slf4j
 @CompileStatic
 class SnowflakePlugin extends BasePlugin {
 
@@ -39,5 +44,8 @@ class SnowflakePlugin extends BasePlugin {
         // Configure Snowflake JDBC to use SLF4J instead of java.util.logging
         // This redirects Snowflake JDBC logs to Nextflow's logging system
         System.setProperty('net.snowflake.jdbc.loggerImpl', 'net.snowflake.client.log.SLF4JLogger')
+
+        // Register the Snowflake FileSystemProvider
+        FileHelper.getOrInstallProvider(SnowflakeFileSystemProvider)
     }
 }

plugins/nf-snowflake/src/main/nextflow/snowflake/SnowflakeTaskHandler.groovy

@@ -156,12 +156,20 @@ class SnowflakeTaskHandler extends TaskHandler {
         this.connection = this.connectionPool.getConnection()
         this.statement = connection.createStatement()
 
+        // Create TaskBean and set defaults
         final TaskBean taskBean = new TaskBean(task)
         if (taskBean.scratch == null) {
             taskBean.scratch = scratchDir
         }
-        final SnowflakeFileCopyStrategy fileCopyStrategy = new SnowflakeFileCopyStrategy(taskBean, executor)
-        final BashWrapperBuilder builder = new BashWrapperBuilder(taskBean, fileCopyStrategy)
+
+        // Always enable stats collection to ensure .command.trace is generated
+        // Nextflow core always attempts to read this file during task finalization
+        taskBean.statsEnabled = true
+        
+        // Use factory method to create wrapper builder (follows Fusion pattern)
+        // This will translate snowflake:// paths to container mount paths
+        // and create the SnowflakeFileCopyStrategy with bin directory support
+        final SnowflakeWrapperBuilder builder = SnowflakeWrapperBuilder.create(taskBean, executor)
         builder.build()
 
         final String spec = buildJobServiceSpec()
@@ -228,10 +236,25 @@ from specification
 
         final String workDir = executor.getWorkDir().toUriString()
 
-        final String workDirStageEnv = System.getenv("workDirStage")
-        final String workDirStage = workDirStageEnv != null ? workDirStageEnv :
-            executor.snowflakeConfig.get("workDirStage")
-        result.addWorkDirMount(workDir, String.format("%s/", workDirStage))
+        // Check if workDir uses snowflake:// scheme
+        if (SnowflakeUri.isSnowflakeStageUri(workDir)) {
+            // Extract stage name using centralized parsing
+            String stageName = SnowflakeUri.extractStageName(workDir)
+
+            // Mount the stage to /mnt/stage/<lowercase_stage_name>
+            String mountPath = "/mnt/stage/${stageName.toLowerCase()}"
+            result.addWorkDirMount(mountPath, stageName)
+
+            log.debug("Mounting Snowflake stage @${stageName} to ${mountPath}")
+        } else {
+            // Legacy behavior: use workDirStage config
+            final String workDirStageEnv = System.getenv("workDirStage")
+            final String workDirStage = workDirStageEnv != null ? workDirStageEnv :
+                executor.snowflakeConfig.get("workDirStage")
+            if (workDirStage) {
+                result.addWorkDirMount(workDir, String.format("%s/", workDirStage))
+            }
+        }
 
         result.addLocalVolume(scratchDir)
 
@@ -312,7 +335,11 @@ from specification
 
     private static List<String> classicSubmitCli(TaskRun task) {
         final result = new ArrayList(BashWrapperBuilder.BASH)
-        result.add("${Escape.path(task.workDir)}/${TaskRun.CMD_RUN}".toString())
+
+        // Translate workDir if it's a snowflake:// path
+        String workDirPath = SnowflakeUri.translateToMount(task.workDir).toUriString()
+        result.add("${Escape.path(workDirPath)}/${TaskRun.CMD_RUN}".toString())
+
         return result
     }
 }