Add support for tower workspaceId [ci fast]

nextflow-io · Mar 19, 2021 · e297768 · e297768
1 parent 6daa21b
commit e297768
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 8 deletions.
diff --git a/plugins/nf-tower/src/main/io/seqera/tower/plugin/TowerClient.groovy b/plugins/nf-tower/src/main/io/seqera/tower/plugin/TowerClient.groovy
@@ -123,6 +123,8 @@ class TowerClient implements TraceObserver {
 
     private String refreshToken
 
+    private String workspaceId 
+
     /**
      * Constructor that consumes a URL and creates
      * a basic HTTP client.
@@ -152,7 +154,7 @@ class TowerClient implements TraceObserver {
     String getRunName() { runName }
 
     String getRunId() { runId }
-    
+
     void setAliveInterval(Duration d) {
         this.aliveInterval = d
     }
@@ -173,6 +175,12 @@ class TowerClient implements TraceObserver {
         this.backOffDelay = value
     }
 
+    void setWorkspaceId( String workspaceId ) {
+        this.workspaceId = workspaceId
+    }
+
+    String getWorkspaceId() { workspaceId }
+
     /**
      * Check the URL and create an HttpPost() object. If a invalid i.e. protocol is used,
      * the constructor will raise an exception.
@@ -197,23 +205,38 @@ class TowerClient implements TraceObserver {
     }
 
     protected String getUrlTraceCreate() {
-        this.endpoint + '/trace/create'
+        def result = this.endpoint + '/trace/create'
+        if( workspaceId )
+            result += "?workspaceId=$workspaceId"
+        return result
     }
 
     protected String getUrlTraceBegin() {
-        "$endpoint/trace/$workflowId/begin"
+        def result = "$endpoint/trace/$workflowId/begin"
+        if( workspaceId )
+            result += "?workspaceId=$workspaceId"
+        return result
     }
 
     protected String getUrlTraceComplete() {
-        "$endpoint/trace/$workflowId/complete"
+        def result = "$endpoint/trace/$workflowId/complete"
+        if( workspaceId )
+            result += "?workspaceId=$workspaceId"
+        return result
     }
 
     protected String getUrlTraceHeartbeat() {
-        "$endpoint/trace/$workflowId/heartbeat"
+        def result = "$endpoint/trace/$workflowId/heartbeat"
+        if( workspaceId )
+            result += "?workspaceId=$workspaceId"
+        return result
     }
 
     protected String getUrlTraceProgress() {
-        "$endpoint/trace/$workflowId/progress"
+        def result = "$endpoint/trace/$workflowId/progress"
+        if( workspaceId )
+            result += "?workspaceId=$workspaceId"
+        return result
     }
 
     /**

diff --git a/plugins/nf-tower/src/main/io/seqera/tower/plugin/TowerFactory.groovy b/plugins/nf-tower/src/main/io/seqera/tower/plugin/TowerFactory.groovy
@@ -27,6 +27,12 @@ import nextflow.util.SimpleHttpClient
 @CompileStatic
 class TowerFactory implements TraceObserverFactory {
 
+    private Map<String,String> env
+
+    TowerFactory(){
+        env = System.getenv()
+    }
+
     @Override
     Collection<TraceObserver> create(Session session) {
         final config = session.config
@@ -39,7 +45,7 @@ class TowerFactory implements TraceObserverFactory {
             return Collections.emptyList()
 
         if ( !endpoint || endpoint=='-' )
-            endpoint = System.getenv('TOWER_API_ENDPOINT') ?: TowerClient.DEF_ENDPOINT_URL
+            endpoint = env.get('TOWER_API_ENDPOINT') ?: TowerClient.DEF_ENDPOINT_URL
 
         final tower = new TowerClient(endpoint)
         if( aliveInterval )
@@ -50,7 +56,7 @@ class TowerFactory implements TraceObserverFactory {
         tower.maxRetries = config.navigate('tower.maxRetries', 5) as int
         tower.backOffBase = config.navigate('tower.backOffBase', SimpleHttpClient.DEFAULT_BACK_OFF_BASE) as int
         tower.backOffDelay = config.navigate('tower.backOffDelay', SimpleHttpClient.DEFAULT_BACK_OFF_DELAY  ) as int
-
+        tower.workspaceId = config.navigate('tower.workspaceId', env.get('TOWER_WORKSPACE_ID'))
         final result = new ArrayList(1)
         result.add(tower)
         return result

diff --git a/plugins/nf-tower/src/test/io/seqera/tower/plugin/TowerClientTest.groovy b/plugins/nf-tower/src/test/io/seqera/tower/plugin/TowerClientTest.groovy
@@ -376,6 +376,20 @@ class TowerClientTest extends Specification {
         tower.getUrlTraceComplete() == 'https://tower.nf/trace/12345/complete'
     }
 
+    def 'should get trace endpoint with workspace' () {
+        given:
+        def tower = new TowerClient('https://tower.nf')
+        tower.workflowId = '12345'
+        tower.workspaceId = '300'
+
+        expect:
+        tower.getUrlTraceCreate() == 'https://tower.nf/trace/create?workspaceId=300'
+        tower.getUrlTraceBegin() == 'https://tower.nf/trace/12345/begin?workspaceId=300'
+        tower.getUrlTraceProgress() == 'https://tower.nf/trace/12345/progress?workspaceId=300'
+        tower.getUrlTraceHeartbeat() == 'https://tower.nf/trace/12345/heartbeat?workspaceId=300'
+        tower.getUrlTraceComplete() == 'https://tower.nf/trace/12345/complete?workspaceId=300'
+    }
+
     def 'should set the auth token' () {
         given:
         def http = Mock(SimpleHttpClient)

diff --git a/plugins/nf-tower/src/test/io/seqera/tower/plugin/TowerFactoryTest.groovy b/plugins/nf-tower/src/test/io/seqera/tower/plugin/TowerFactoryTest.groovy
@@ -57,4 +57,25 @@ class TowerFactoryTest extends Specification {
 
     }
 
+    def 'should create with with workspace id'() {
+        given:
+        def session = Mock(Session)
+        def factory = new TowerFactory(env: [TOWER_WORKSPACE_ID: '100'])
+
+        when:
+        def client = (TowerClient) factory.create(session)[0]
+        then:
+        session.getConfig() >> [tower: [enabled: true]]
+        and:
+        client.getWorkspaceId() == '100'
+
+
+        when:
+        client = (TowerClient) factory.create(session)[0]
+        then:
+        session.getConfig() >> [tower: [enabled: true, workspaceId: '200']]
+        and:
+        client.getWorkspaceId() == '200'
+    }
+
 }