##################################################################
# Copyright 2018 Open Source Geospatial Foundation and others    #
# licensed under MIT, Please consult LICENSE.txt for details     #
##################################################################

import os
import pywps.configuration as config
from pywps.processing.basic import Processing
from pywps.exceptions import SchedulerNotAvailable
from pywps.response.status import WPS_STATUS

import logging
LOGGER = logging.getLogger("PYWPS")


class Scheduler(Processing):
    """
    :class:`Scheduler` is processing implementation to run jobs on schedulers
    like slurm, grid-engine and torque. It uses the drmaa python library
    as client to launch jobs on a scheduler system.

    See: http://drmaa-python.readthedocs.io/en/latest/index.html
    """

    def start(self):
        self.job.wps_response._update_status(WPS_STATUS.ACCEPTED, 'Submitting job ...', 0)
        # run remote pywps process
        jobid = self.run_job()
        self.job.wps_response._update_status(WPS_STATUS.ACCEPTED,
                                             'Your job has been submitted with ID {}'.format(jobid), 0)

    def run_job(self):
        LOGGER.info("Submitting job ...")
        try:
            import drmaa
            session = drmaa.Session()
            # init session
            session.initialize()
            # dump job to file
            dump_filename = self.job.dump()
            if not dump_filename:
                raise Exception("Could not dump job status.")
            # prepare remote command
            jt = session.createJobTemplate()
            jt.remoteCommand = os.path.join(
                config.get_config_value('processing', 'path'),
                'joblauncher')
            if os.getenv("PYWPS_CFG"):
                import shutil
                cfg_file = os.path.join(self.job.workdir, "pywps.cfg")
                shutil.copy2(os.getenv('PYWPS_CFG'), cfg_file)
                LOGGER.debug("Copied pywps config: {}".format(cfg_file))
                jt.args = ['-c', cfg_file, dump_filename]
            else:
                jt.args = [dump_filename]
            drmaa_native_specification = config.get_config_value('processing', 'drmaa_native_specification')
            if drmaa_native_specification:
                jt.nativeSpecification = drmaa_native_specification
            jt.joinFiles = True
            jt.outputPath = ":{}".format(os.path.join(self.job.workdir, "job-output.txt"))
            # run job
            jobid = session.runJob(jt)
            LOGGER.info('Your job has been submitted with ID {}'.format(jobid))
            # show status
            import time
            time.sleep(1)
            LOGGER.info('Job status: {}'.format(session.jobStatus(jobid)))
            # Cleaning up
            session.deleteJobTemplate(jt)
            # close session
            session.exit()
        except Exception as e:
            raise SchedulerNotAvailable("Could not submit job: {}".format(str(e)))
        return jobid

@@ -1288,17 +1288,12 @@
     
-session =
+with
  drm
@@ -1308,35 +1308,19 @@
 on()
-%0A            # init
+ as
  session
 %0A   
@@ -1319,42 +1319,16 @@
 sion
+:
 %0A      
-      session.initialize()%0A
     
@@ -1333,18 +1333,16 @@
         
-  
 # dump j
@@ -1352,16 +1352,20 @@
 to file%0A
+    
         
@@ -1412,16 +1412,20 @@
         
+  
+  
 if not d
@@ -1438,16 +1438,20 @@
 lename:%0A
+    
         
@@ -1516,16 +1516,20 @@
         
+ 
+   
 # prepar
@@ -1545,16 +1545,20 @@
 command%0A
+    
         
@@ -1606,16 +1606,20 @@
         
+ 
+   
 jt.remot
@@ -1643,16 +1643,20 @@
 h.join(%0A
+    
         
@@ -1726,16 +1726,20 @@
         
+    
 'joblaun
@@ -1745,16 +1745,20 @@
 ncher')%0A
+    
         
@@ -1800,24 +1800,28 @@
             
+    
 import shuti
@@ -1838,16 +1838,20 @@
         
+   
+ 
 cfg_file
@@ -1913,16 +1913,20 @@
         
+    
 shutil.c
@@ -1976,24 +1976,28 @@
             
+  
+  
 LOGGER.debug
@@ -2049,32 +2049,36 @@
                 
+    
 jt.args = %5B'-c',
@@ -2119,14 +2119,22 @@
     
+ 
+   
 else:%0A
+    
     
@@ -2163,32 +2163,36 @@
 %5Bdump_filename%5D%0A
+    
             drma
@@ -2284,16 +2284,20 @@
 ation')%0A
+    
         
@@ -2347,16 +2347,20 @@
         
+    
 jt.nativ
@@ -2399,24 +2399,28 @@
 ecification%0A
+    
             
@@ -2451,16 +2451,20 @@
         
+    
 jt.outpu
@@ -2542,24 +2542,28 @@
             
+   
+ 
 # run job%0A  
@@ -2560,16 +2560,20 @@
 run job%0A
+    
         
@@ -2599,24 +2599,28 @@
 .runJob(jt)%0A
+    
             
@@ -2699,16 +2699,20 @@
         
+    
 # show s
@@ -2725,24 +2725,28 @@
             
+  
+  
 import time%0A
@@ -2745,16 +2745,20 @@
 rt time%0A
+    
         
@@ -2775,16 +2775,20 @@
 leep(1)%0A
+    
         
@@ -2850,24 +2850,28 @@
 us(jobid)))%0A
+    
             
@@ -2888,32 +2888,36 @@
  up%0A            
+    
 session.deleteJo
@@ -2934,63 +2934,8 @@
 jt)%0A
-            # close session%0A            session.exit()%0A
     

use with statement to close drmaa session