Delay loading of some modules (#690)

* Delay loading of modules

* remove old function call

* Add command to config

turbinia/config/__init__.py

@@ -92,7 +92,9 @@
     # Prometheus config
     'PROMETHEUS_ENABLED',
     'PROMETHEUS_ADDR',
-    'PROMETHEUS_PORT'
+    'PROMETHEUS_PORT',
+    # General config
+    'TURBINIA_COMMAND'
 ]
 
 # Environment variable to look for path data in

turbinia/turbiniactl.py

@@ -29,8 +29,6 @@
 from turbinia import config
 from turbinia import TurbiniaException
 from turbinia.config import logger
-from libcloudforensics.providers.gcp import forensics as gcp_forensics
-from turbinia.lib import google_cloud
 from turbinia import __version__
 from turbinia.processors import archive
 from turbinia.output_manager import OutputManager
@@ -430,6 +428,8 @@ def main():
     google_cloud.setup_stackdriver_handler(
         config.TURBINIA_PROJECT, args.command)
 
+  config.TURBINIA_COMMAND = args.command
+
   log.info('Turbinia version: {0:s}'.format(__version__))
 
   # Do late import of other needed Turbinia modules.  This is needed because the
@@ -445,6 +445,10 @@ def main():
   from turbinia import evidence
   from turbinia.message import TurbiniaRequest
 
+  if config.TASK_MANAGER == 'PSQ':
+    from turbinia.lib import google_cloud
+    from libcloudforensics.providers.gcp import forensics as gcp_forensics
+
   # Print out config if requested
   if args.command == 'config':
     if args.file_only:

turbinia/workers/__init__.py

@@ -440,6 +440,22 @@ def deserialize(cls, input_dict):
         input_dict['last_update'], DATETIME_FORMAT)
     return task
 
+  @classmethod
+  def check_worker_role(cls):
+    """Checks whether the execution context is within a worker or nosetests.
+
+    Returns:
+      bool: If the current execution is in a worker or nosetests.
+    """
+    if config.TURBINIA_COMMAND in ('celeryworker', 'psqworker'):
+      return True
+
+    for arg in sys.argv:
+      if 'nosetests' in arg:
+        return True
+
+    return False
+
   def evidence_setup(self, evidence):
     """Validates and processes the evidence.
 

turbinia/workers/partitions.py

@@ -14,18 +14,24 @@
 # limitations under the License.
 """Task for enumerating partitions in a disk."""
 
-from dfvfs.helpers import volume_scanner
-from dfvfs.lib import definitions as dfvfs_definitions
-from dfvfs.lib import errors as dfvfs_errors
-from dfvfs.volume import tsk_volume_system
-
 from turbinia import TurbiniaException
 from turbinia.evidence import RawDiskPartition
-from turbinia.lib import dfvfs_classes
 from turbinia.lib import text_formatter as fmt
 from turbinia.workers import Priority
 from turbinia.workers import TurbiniaTask
 
+if TurbiniaTask.check_worker_role():
+  try:
+    from dfvfs.helpers import volume_scanner
+    from dfvfs.lib import definitions as dfvfs_definitions
+    from dfvfs.lib import errors as dfvfs_errors
+    from dfvfs.volume import tsk_volume_system
+
+    from turbinia.lib import dfvfs_classes
+  except ImportError as exception:
+    message = 'Could not import dfVFS libraries: {0!s}'.format(exception)
+    raise TurbiniaException(message)
+
 
 class PartitionEnumerationTask(TurbiniaTask):
   """Task to enumerate partitions in a disk."""