def test_invalid_object_path(self, mock_parent_init): # This is just the path of a bucket hence invalid filename file_name = 'gs://test-bucket' mock_parent_init.return_value = None gcs_bucket_helper = GoogleCloudBucketHelper() gcs_bucket_helper._gcs_hook = mock.Mock() with self.assertRaises(Exception) as context: gcs_bucket_helper.google_cloud_to_local(file_name) self.assertEqual( 'Invalid Google Cloud Storage (GCS) object path: {}'.format(file_name), str(context.exception))
def test_invalid_object_path(self, mock_parent_init): # This is just the path of a bucket hence invalid filename file_name = 'gs://test-bucket' mock_parent_init.return_value = None gcs_bucket_helper = GoogleCloudBucketHelper() gcs_bucket_helper._gcs_hook = mock.Mock() with self.assertRaises(Exception) as context: gcs_bucket_helper.google_cloud_to_local(file_name) self.assertEqual( 'Invalid Google Cloud Storage (GCS) object path: {}'.format( file_name), str(context.exception))
def execute(self, context): bucket_helper = GoogleCloudBucketHelper(self.gcp_conn_id, self.delegate_to) self.py_file = bucket_helper.google_cloud_to_local(self.py_file) hook = GFWDataFlowHook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to, poll_sleep=self.poll_sleep) dataflow_options = self.dataflow_default_options.copy() dataflow_options.update(self.options) camel_to_snake = lambda name: re.sub( r'[A-Z]', lambda x: '_' + x.group(0).lower(), name) formatted_options = { camel_to_snake(key): dataflow_options[key] for key in dataflow_options } hook.start_python_dataflow(self.job_name, formatted_options, self.py_file, self.py_options)
def test_valid_object(self, mock_parent_init): file_name = 'gs://test-bucket/path/to/obj.jar' mock_parent_init.return_value = None gcs_bucket_helper = GoogleCloudBucketHelper() gcs_bucket_helper._gcs_hook = mock.Mock() def _mock_download(bucket, object, filename=None): text_file_contents = 'text file contents' with open(filename, 'w') as text_file: text_file.write(text_file_contents) return text_file_contents gcs_bucket_helper._gcs_hook.download.side_effect = _mock_download local_file = gcs_bucket_helper.google_cloud_to_local(file_name) self.assertIn('obj.jar', local_file)
def test_valid_object(self, mock_parent_init): file_name = 'gs://test-bucket/path/to/obj.jar' mock_parent_init.return_value = None gcs_bucket_helper = GoogleCloudBucketHelper() gcs_bucket_helper._gcs_hook = mock.Mock() def _mock_download(bucket, object, filename=None): text_file_contents = 'text file contents' with open(filename, 'w') as text_file: text_file.write(text_file_contents) return text_file_contents gcs_bucket_helper._gcs_hook.download.side_effect = _mock_download local_file = gcs_bucket_helper.google_cloud_to_local(file_name) self.assertIn('obj.jar', local_file)
def execute(self, context): """Execute the python dataflow job.""" bucket_helper = GoogleCloudBucketHelper( self.gcp_conn_id, self.delegate_to) self.py_file = bucket_helper.google_cloud_to_local(self.py_file) hook = DataFlow3Hook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to, poll_sleep=self.poll_sleep) dataflow_options = self.dataflow_default_options.copy() dataflow_options.update(self.options) # Convert argument names from lowerCamelCase to snake case. camel_to_snake = lambda name: re.sub( r'[A-Z]', lambda x: '_' + x.group(0).lower(), name) formatted_options = {camel_to_snake(key): dataflow_options[key] for key in dataflow_options} hook.start_python_dataflow( self.job_name, formatted_options, self.py_file, self.py_options, py_interpreter="python3")
def execute_direct_runner(self, context): bucket_helper = GoogleCloudBucketHelper(self.gcp_conn_id, self.delegate_to) self.py_file = bucket_helper.google_cloud_to_local(self.py_file) dataflow_options = self.dataflow_default_options.copy() dataflow_options.update(self.options) # Convert argument names from lowerCamelCase to snake case. camel_to_snake = lambda name: re.sub( r'[A-Z]', lambda x: '_' + x.group(0).lower(), name) formatted_options = { camel_to_snake(key): dataflow_options[key] for key in dataflow_options } hook = DataFlowDirectRunnerHook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to) hook.start_python_dataflow(self.task_id, formatted_options, self.py_file, self.py_options) pass
def test_empty_object(self, mock_parent_init): file_name = 'gs://test-bucket/path/to/obj.jar' mock_parent_init.return_value = None gcs_bucket_helper = GoogleCloudBucketHelper() gcs_bucket_helper._gcs_hook = mock.Mock() def _mock_download(bucket, object, filename=None): text_file_contents = '' with open(filename, 'w') as text_file: text_file.write(text_file_contents) return text_file_contents gcs_bucket_helper._gcs_hook.download.side_effect = _mock_download with self.assertRaises(Exception) as context: gcs_bucket_helper.google_cloud_to_local(file_name) self.assertEqual( 'Failed to download Google Cloud Storage (GCS) object: {}'.format(file_name), str(context.exception))
def test_empty_object(self, mock_parent_init): file_name = 'gs://test-bucket/path/to/obj.jar' mock_parent_init.return_value = None gcs_bucket_helper = GoogleCloudBucketHelper() gcs_bucket_helper._gcs_hook = mock.Mock() def _mock_download(bucket, object, filename=None): text_file_contents = '' with open(filename, 'w') as text_file: text_file.write(text_file_contents) return text_file_contents gcs_bucket_helper._gcs_hook.download.side_effect = _mock_download with self.assertRaises(Exception) as context: gcs_bucket_helper.google_cloud_to_local(file_name) self.assertEqual( 'Failed to download Google Cloud Storage (GCS) object: {}'.format( file_name), str(context.exception))
def execute(self, context): try: self.log.info("| correlationId={0} | op=jar-execution | status=OK | desc=Started processing task '{1}'". format(self.correlation_id, self.task_id)) bucket_helper = GoogleCloudBucketHelper(self.gcp_conn_id, self.delegate_to) self.jar = bucket_helper.google_cloud_to_local(self.jar) command = 'java -jar {0}'.format(self.jar) if self.options is not None: for attr, value in self.options.iteritems(): command += " -" + attr + " \"" + value + "\"" self.bash_command = command super(JVMOperator, self).execute(context) self.log.info("| correlationId={0} | op=jar-execution | status=OK | desc=Completed processing task '{1}'". format(self.correlation_id, self.task_id)) except Exception as exception: self.log.error( "| correlationId={0} | op=jar-execution | status=KO | desc=Failed to process the task '{1}' with exception '{2}'". format(self.correlation_id, self.task_id, exception)) raise exception