TobKed commented on a change in pull request #8553:
URL: https://github.com/apache/airflow/pull/8553#discussion_r501799152



##########
File path: airflow/providers/google/cloud/hooks/dataflow.py
##########
@@ -783,6 +794,77 @@ def cancel_job(
             name=job_name,
             job_id=job_id,
             location=location,
-            poll_sleep=self.poll_sleep
+            poll_sleep=self.poll_sleep,
+            num_retries=self.num_retries,
         )
         jobs_controller.cancel()
+
+    @GoogleBaseHook.fallback_to_default_project_id
+    def start_sql_job(
+        self,
+        job_name: str,
+        query: str,
+        options: Dict[str, Any],
+        project_id: str,
+        location: str = DEFAULT_DATAFLOW_LOCATION,
+        on_new_job_id_callback: Optional[Callable[[str], None]] = None
+    ):
+        """
+        Starts Dataflow SQL query.
+
+        :param job_name: The unique name to assign to the Cloud Dataflow job.
+        :type job_name: str
+        :param query: The SQL query to execute.
+        :type query: str
+        :param options: Job parameters to be executed.
+            For more information, look at:
+            
`https://cloud.google.com/sdk/gcloud/reference/beta/dataflow/sql/query
+            <gcloud beta dataflow sql query>`__
+            command reference
+        :param location: The location of the Dataflow job (for example 
europe-west1)
+        :type location: str
+        :param project_id: The ID of the GCP project that owns the job.
+            If set to ``None`` or missing, the default project_id from the GCP 
connection is used.
+        :type project_id: Optional[str]
+        :param on_new_job_id_callback: Callback called when the job ID is 
known.
+        :type on_new_job_id_callback: callable
+        :return: the new job object
+        """
+        cmd = [
+            'gcloud',
+            'beta',
+            'dataflow',
+            'sql',
+            'query',
+            query,
+            f'--project={project_id}',
+            '--format=value(job.id)',
+            f'--job-name={job_name}',
+            f'--region={location}',
+            *(self._options_to_args(options))
+        ]
+        self.log.info("Executing command: %s", " ".join([shlex.quote(c) for c 
in cmd]))
+        with self.provide_authorized_gcloud():
+            proc = subprocess.run(  # pylint: disable=subprocess-run-check
+                cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        self.log.info("Output: %s", proc.stdout.decode())
+        self.log.info("Stderr: %s", proc.stderr.decode())

Review comment:
       I changed it :)




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


Reply via email to