You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by GitBox <gi...@apache.org> on 2020/07/19 02:05:37 UTC
[GitHub] [airflow] ephraimbuddy opened a new pull request #9879: Implement Compute Engine SSH hook and operator
ephraimbuddy opened a new pull request #9879:
URL: https://github.com/apache/airflow/pull/9879
---
Closes #9752
- [ ] Add Compute Engine SSH Hook and Operator
- [ ] Add Example DAG and Documentation
- [ ] Add System tests
Make sure to mark the boxes below before creating PR: [x]
- [x] Description above provides context of the change
- [x] Unit tests coverage for changes (not needed for documentation changes)
- [x] Target Github ISSUE in description if exists
- [x] Commits follow "[How to write a good git commit message](http://chris.beams.io/posts/git-commit/)"
- [x] Relevant documentation is updated including usage instructions.
- [x] I will engage committers as explained in [Contribution Workflow Example](https://github.com/apache/airflow/blob/master/CONTRIBUTING.rst#contribution-workflow-example).
---
In case of fundamental code change, Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvements+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in [UPDATING.md](https://github.com/apache/airflow/blob/master/UPDATING.md).
Read the [Pull Request Guidelines](https://github.com/apache/airflow/blob/master/CONTRIBUTING.rst#pull-request-guidelines) for more information.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713213203
[The Workflow run](https://github.com/apache/airflow/actions/runs/318795591) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713555164
[The Workflow run](https://github.com/apache/airflow/actions/runs/319871248) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-712822580
Hi.
I worked on this change today/yesterday and I have good news.
- Authentication with gcloud works. Previously, there were problems with the lack of an email address.
- IAP tunneling works
I still need to work on authorization with OSLogin.
Best regards,
Kamil Breguła
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] potiuk commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
potiuk commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713470547
There is apparently a hanging check - that cannot be cancelled even manually with this one. We could close the PR but might be worth opening a support ticket to GitHub @mik-laj as discussed.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713323233
[The Workflow run](https://github.com/apache/airflow/actions/runs/319212260) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] ephraimbuddy commented on a change in pull request #9879: [WIP]Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
ephraimbuddy commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r458613130
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,250 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import subprocess
+import time
+import uuid
+from typing import Any, List, Optional, Union
+
+import paramiko
+from googleapiclient.discovery import build
+from paramiko import SSHException
+
+from airflow import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+
+class ComputeEngineSshHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version: The Oslogin discovery API version
+ :type api_version: str
+ :param gcp_conn_id: The connection id to use when fetching connection info
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+
+ def __init__(self,
+ api_version: str = "v1",
+ gcp_conn_id: str = "google_cloud_default",
+ delegate_to: Optional[str] = None
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+ self.conn = None
+
+ def get_conn(self) -> Any:
+ """
+ Connects to Oslogin API
+ """
+ if self.conn is None:
+ http_authorized = self._authorize()
+ self.conn = build("oslogin", self.api_version,
+ http=http_authorized, cache_discovery=False)
+ return self.conn
+
+ def get_sshclient(self,
+ private_key_file: str,
+ hostname: str,
+ username: str,
+ strict_host_checking: bool = True
+ ) -> paramiko.SSHClient:
+ """
+ Get SSHClient
+
+ :param private_key_file: The filename of the private key
+ :type private_key_file: str
+ :param hostname: The hostname of the target instance
+ :type hostname: str
+ :param username: The POXIS username from the login profile of the user account
+ :type username: str
+ :param strict_host_checking: Whether to verify key
+ :type strict_host_checking: bool
+ """
+
+ self.log.info("Creating remote connection to host")
+ client = paramiko.SSHClient()
+ if not strict_host_checking:
+ self.log.warning('No Host Key Verification. This wont protect '
+ 'against Man-In-The-Middle attacks')
+ # Default is RejectPolicy
+ client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+ client.connect(hostname=hostname,
+ username=username,
+ key_filename=private_key_file)
+ return client
+
+ def get_service_account(self) -> str:
+ """
+ Get the service account id
+ """
+ account: str = self._get_credentials().service_account_email
+ if not account.startswith('users/'):
+ account = 'users/' + account
+ return account
+
+ def get_username(self) -> str:
+ """
+ Using the OS Login API, get the POSIX username from the login profile
+ for the service account.
+
+ """
+ conn = self.get_conn()
+ account = self.get_service_account()
+ profile = conn.users().getLoginProfile(name=account).execute_in_compute() # pylint: disable=no-member
+ username: str = profile.get('posixAccounts')[0].get('username')
+ return username
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_hostname(self, project_id: Optional[str], instance: str, zone: str,
+ hostname: Optional[str] = None) -> str:
+ """
+ Create the target instance hostname
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param hostname: The hostname of the target instance. If not provided,
+ it is generated from zone, instance name and project id
+ :type hostname: str
+ """
+ if not hostname:
+ hostname = f'{instance}.{zone}.c.{project_id}.internal'
+ return hostname
+
+ def execute_local(self, cmd: Union[str, List[str]]) -> None:
+ """
+ Executes command in local machine
+
+ :param cmd: The command to execute
+ :type cmd: Union[str, List[str]]
+ """
+ self.log.info('Executing command: %s', str(cmd))
+ process = subprocess.Popen(cmd)
+ output = process.communicate()[0]
+ return_code = process.returncode
+ if return_code:
+ # Error
+ raise AirflowException(f'The Command {cmd} returned error status {return_code}')
+ if output:
+ self.log.info(output)
+
+ def create_ssh_key(self,
+ account: str,
+ expire_time: int = 300) -> str:
+ """
+ Generate a temporary SSH key and apply it to the specified account
+
+ :param account: The service account id
+ :type account: str
+ :param expire_time: The maximum amount of time before the private key is expired
+ :type expire_time: int
+ """
+ try:
+ self.log.info("Generating ssh keys...")
+ private_key_file = '/tmp/key-' + str(uuid.uuid4())
+ private_key_obj = paramiko.RSAKey.generate(2048)
+ private_key_obj.write_private_key_file(private_key_file)
+ except (IOError, SSHException) as err:
+ raise AirflowException(f"Error encountered creating ssh keys, {err}")
+
+ with open(private_key_file + '.pub', 'r') as original:
+ public_key = original.read().strip()
+
+ # Expiration time is in microseconds.
+ expiration = int((time.time() + expire_time) * 1000000)
+
+ body = {
+ 'key': public_key,
+ 'expirationTimeUsec': expiration,
+ }
+ self.get_conn(). \
+ users().importSshPublicKey( # pylint: disable=no-member
+ parent=account, body=body).execute_in_compute()
+ return private_key_file
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def run(self,
+ cmd: str, instance: str, zone: str,
+ project_id: Optional[str] = None,
+ expire_time: int = 300,
+ hostname: Optional[str] = None,
+ strict_host_key_checking: bool = True
+ ) -> Any:
+ """
+ Execute a command in a remote instance
+
+ :param cmd: The command to execute on the instance
+ :type cmd: str
+ :param instance: The name of the instance the command will be executed on
+ :type instance: str
+ :param zone: The zone where the instance is located
+ :type zone: str
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param expire_time: The maximum amount of time before the private key is expired
+ :type expire_time: int
+ :param hostname: The hostname of the instance.
+ :type hostname: str
+ :param strict_host_key_checking: Whether to verify the host key
+ :type strict_host_key_checking: bool
+ """
+ hostname = self.get_hostname(project_id=project_id,
+ instance=instance,
+ zone=zone,
+ hostname=hostname)
+
+ account = self.get_service_account()
+
+ username = self.get_username()
+
+ private_key_file = self.create_ssh_key(account=account,
Review comment:
Using project or instance metadata seem discouraged. See
https://cloud.google.com/compute/docs/instances/connecting-advanced#provide-key
What do you think?
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-712868780
[The Workflow run](https://github.com/apache/airflow/actions/runs/317688159) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] ephraimbuddy commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
ephraimbuddy commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-717848935
> I thoroughly rewritten the implementation to be compatible with the SSH Operator and also to make the logic much simpler. I think the code is perfect now. I managed to solve all the problems that I encountered and the system tests are passing correctly. Especially SSH tunneling was problematic as SSH tunnel runs in an independent subprocess and I didn't see any error messages so I didn't know you should get authorization separately but now everything works.
>
> I still have a documentation update + some unit tests.
>
> @ephraimbuddy What do you think about the current implementation? Is it clearer?
Wow! This is much better. I gave up on this after so many attempts. Thanks a bunch!.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713741546
[The Workflow run](https://github.com/apache/airflow/actions/runs/320312735) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] ephraimbuddy removed a comment on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
ephraimbuddy removed a comment on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-665318597
Hi @mik-laj, I have tried everything I could but have not been able to use tunnelling. It seems complicated more than I thought.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-715626059
[The Workflow run](https://github.com/apache/airflow/actions/runs/325094834) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on a change in pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r460596997
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,378 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import subprocess
+import time
+from io import StringIO
+from typing import Any, List, Optional, Union
+
+import paramiko
+from google.cloud.oslogin_v1 import OsLoginServiceClient
+from googleapiclient.discovery import build
+from paramiko import AuthenticationException, BadHostKeyException, SSHException
+
+from airflow import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+
+class ComputeEngineSshHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version: The Compute Engine discovery API version
+ :type api_version: str
+ :param gcp_conn_id: The connection id to use when fetching connection information
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+ _conn: Optional[Any] = None
+
+ def __init__(
+ self,
+ api_version: str = 'v1',
+ gcp_conn_id: str = 'google_cloud_default',
+ delegate_to: Optional[str] = None
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+ self._oslogin_conn: Optional[OsLoginServiceClient] = None
+
+ def get_conn(self) -> Any:
+ """
+ Retrieves connection to Google Compute Engine.
+
+ :return: Google Compute Engine services object
+ :rtype: dict
+ """
+ if not self._conn:
+ http_authorized = self._authorize()
+ self._conn = build('compute', self.api_version,
+ http=http_authorized, cache_discovery=False)
+ return self._conn
+
+ def get_oslogin_conn(self) -> OsLoginServiceClient:
+ """
+ Connects to Oslogin API
+ """
+ if not self._oslogin_conn:
+ self._oslogin_conn = OsLoginServiceClient(
+ credentials=self._get_credentials(), client_info=self.client_info
+ )
+
+ return self._oslogin_conn
+
+ def get_sshclient(self,
+ pkey: str,
+ hostname: str,
+ username: str
+ ) -> paramiko.SSHClient:
+ """
+ Get SSHClient
+
+ :param pkey: The private key
+ :type pkey: str
+ :param hostname: The hostname of the target instance
+ :type hostname: str
+ :param username: The username from the login profile of the user account
+ :type username: str
+ """
+
+ client = paramiko.SSHClient()
+ # Default is RejectPolicy
+ # No knownhost checking since we are not storing pkey
+ client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+ client.connect(hostname=hostname,
+ username=username,
+ pkey=pkey,
+ look_for_keys=False)
+ return client
+
+ def get_service_account(self) -> str:
+ """
+ Get the service account id
+ """
+ account: str = self._get_credentials().service_account_email
+ os_login_client = self.get_oslogin_conn()
+ return os_login_client.user_path(user=account)
+
+ def get_username(self, oslogin=True) -> str:
+ """
+ Get the account username
+
+ :param oslogin: Whether to use OsLogin
+ :type oslogin: bool
+ """
+ account = self._get_credentials().service_account_email
+ if oslogin:
+ profile = self.get_oslogin_conn().get_login_profile(name=self.get_service_account())
+ account = profile.posix_accounts[0]
+ return account.username
+ return account.split("@")[0]
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_connection_type(self, project_id: Optional[str], instance: str, zone: str,
+ hostname: Optional[str] = None,
+ use_iap_tunnel: bool = False,
+ use_internal_ip: bool = False
+ ) -> str:
+ """
+ Selects how to connect to the target instance
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param hostname: The hostname of the target instance. If provided, this would
+ be used, regardless of the setting for use_iap_tunnel and use_internal_ip
+ :type hostname: str
+ :param use_iap_tunnel: Whether to connect through IAP tunnel
+ :type use_iap_tunnel: bool
+ :param use_internal_ip: Whether to connect using internal IP
+ :type use_internal_ip: bool
+ """
+ if not hostname:
+ instance_info = self.instance_info(project_id=project_id,
+ zone=zone,
+ instance=instance)
+ if use_iap_tunnel or use_internal_ip:
+ return instance_info["networkInterfaces"][0] \
+ .get("networkIP")
+ else:
+ access_config = instance_info \
+ .get("networkInterfaces")[0] \
+ .get("accessConfigs")
+ if access_config:
+ return access_config[0].get("natIP")
+ else:
+ raise AirflowException("The target instance does not have external IP,"
+ "consider specifying use_internal_ip or use_iap_tunnel")
+ return hostname
+
+ def execute(self, cmd: Union[str, List[str]]) -> None:
+ """
+ Executes command in local machine
+
+ :param cmd: The command to execute
+ :type cmd: Union[str, List[str]]
+ """
+ self.log.info('Executing command: %s', str(cmd))
+ process = subprocess.Popen(cmd)
+ output = process.communicate()[0]
+ returncode = process.returncode
+ if returncode:
+ # Error
+ self.log.error('Command returned error status %s', returncode)
+ if output:
+ self.log.info(output)
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def instance_info(self, project_id, zone, instance):
+ """
+ Used to get instance information
+
+ :param project_id: The project ID where the instance is located
+ :type project_id: str
+ :param zone: The zone of the instance
+ :type zone: str
+ :param instance: The name of the instance
+ :type instance: str
+ """
+ instance_info = self.get_conn().instances().get( # pylint: disable=no-member
+ project=project_id,
+ instance=instance,
+ zone=zone).execute()
+ return instance_info
+
+ def _update_metadata(self, keys, metadata):
+ """
+ Updates ssh keys on metadata dictionary
+
+ :param keys: The public key to set on the metadata
+ :type keys: str
+ :param metadata: A dictionary containing instance metadata
+ :type metadata: dict
+ """
+ items = metadata.get("items")
+ for item in items:
+ if item.get("key") == "ssh-keys":
+ keys += item["value"]
+ item['value'] = keys
+ break
+ else:
+ new_dict = dict(key='ssh-keys', value=keys)
+ metadata['items'] = [new_dict]
+ return metadata
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def set_sshkey_metadata(self, project_id, zone, instance, username, pubkey):
+ """
+ Sets ssh key on instance metadata
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param username: The logged in account username
+ :type username: str
+ :param pubkey: The public key to set on the instance
+ :type pubkey: str
+ """
+ instance_info = self.instance_info(project_id=project_id,
+ zone=zone,
+ instance=instance)
+ keys = username + ":" + pubkey + "\n"
+ metadata = instance_info['metadata']
+ metadata = self._update_metadata(keys, metadata)
+ self.get_conn().instances().setMetadata( # pylint: disable=no-member
+ project=project_id,
+ zone=zone,
+ instance=instance,
+ body=metadata
+ ).execute()
+ time.sleep(10) # TODO improve on this. The status should be Done before proceeding
+
+ def set_oslogin_key(self, pubkey, expire_time):
+ """
+ Sets public key on instance os login
+
+ :param pubkey: The public key to set on the OsLogin
+ :type pubkey: str
+ :param expire_time: The amount of time before the key expires in seconds
+ :type expire_time: int
+ """
+ expiration = int((time.time() + expire_time) * 1000000)
+ body = {
+ "key": pubkey,
+ "expiration_time_usec": expiration
+ }
+ os_login_client = self.get_oslogin_conn()
+ os_login_client.import_ssh_public_key(
+ parent=self.get_service_account(),
+ ssh_public_key=body
+ )
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def _create_ssh_key(self,
+ instance: str,
+ zone: str,
+ project_id: Optional[str] = None,
+ expire_time: int = 300,
+ use_oslogin: bool = True) -> Any:
+ """
+ Generate a temporary SSH key and apply it to the specified account
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param expire_time: The maximum amount of time before the private key is expired
+ :type expire_time: int
+ :param use_oslogin: Whether to use OsLogin to manage key
+ :type use_oslogin: int
+ """
+ username = self.get_username(oslogin=use_oslogin)
+ try:
+ self.log.info("Generating ssh keys...")
+ pkey_file = StringIO()
+ pkey_obj = paramiko.RSAKey.generate(2048)
+ pkey_obj.write_private_key(pkey_file)
+ pubkey = f"{pkey_obj.get_name()} {pkey_obj.get_base64()} {username}"
+ except (IOError, SSHException) as err:
+ raise AirflowException(f"Error encountered creating ssh keys, {err}")
+
+ if use_oslogin:
+ # Expiration time is in microseconds.
+ self.set_oslogin_key(pubkey, expire_time)
+ else:
+ # using instance metadata to manage keys
+ self.set_sshkey_metadata(project_id=project_id,
+ zone=zone,
+ instance=instance,
+ username=username,
+ pubkey=pubkey)
+ return pkey_obj
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def run(self,
+ cmd: str, instance: str, zone: str,
+ project_id: Optional[str] = None,
+ expire_time: int = 300,
+ hostname: Optional[str] = None,
+ use_internal_ip: bool = False,
+ use_iap_tunnel: bool = False,
+ use_oslogin: bool = True
+ ) -> Any:
+ """
+ Execute a command in a remote instance
+
+ :param cmd: The command to execute on the instance
+ :type cmd: str
+ :param instance: The name of the instance the command will be executed on
+ :type instance: str
+ :param zone: The zone where the instance is located
+ :type zone: str
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param expire_time: The maximum amount of time before the private key is expired
+ :type expire_time: int
+ :param hostname: The hostname of the instance.
+ :type hostname: str
+ :param use_iap_tunnel: Whether to connect through IAP tunnel
+ :type use_iap_tunnel: bool
+ :param use_internal_ip: Whether to connect using internal IP
+ :type use_internal_ip: bool
+ :param use_oslogin: Whether to manage keys using OsLogin API. If false,
+ keys are managed using instance metadata
+ """
+ hostname = self.get_connection_type(project_id=project_id,
+ instance=instance,
+ zone=zone,
+ hostname=hostname,
+ use_iap_tunnel=use_iap_tunnel,
+ use_internal_ip=use_internal_ip
+ )
+
+ username = self.get_username(oslogin=use_oslogin)
+
+ self.log.info("Creating remote connection to host")
+ pkey = self._create_ssh_key(
+ instance=instance,
+ project_id=project_id,
+ zone=zone,
+ expire_time=expire_time,
+ use_oslogin=use_oslogin
+ )
+ sshclient = self.get_sshclient(pkey=pkey,
Review comment:
There is one problem if internal IP addresses are used. Two servers can use the same IP address if they are on two different networks. We should make sure that we have connected to the expected server.
```
metadata_id_url = 'http://metadata.google.internal/computeMetadata/v1/instance/id'
remote_command = f'[ `curl "{metadata_id_url}" -H "Metadata-Flavor: Google" -q` = {instance_id} ] || exit 23'
```
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-712907350
[The Workflow run](https://github.com/apache/airflow/actions/runs/317812331) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713775094
[The Workflow run](https://github.com/apache/airflow/actions/runs/320410020) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713210187
[The Workflow run](https://github.com/apache/airflow/actions/runs/318782789) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713023568
[The Workflow run](https://github.com/apache/airflow/actions/runs/318143633) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] ephraimbuddy commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
ephraimbuddy commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713012817
> Hi.
> I worked on this change today/yesterday and I have good news.
> - Authentication with gcloud works. Previously, there were problems with the lack of an email address.
> - IAP tunneling works
> I still need to work on authorization with OSLogin.
>
> Best regards,
> Kamil Breguła
Great!
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713230222
[The Workflow run](https://github.com/apache/airflow/actions/runs/318872381) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713018754
[The Workflow run](https://github.com/apache/airflow/actions/runs/318064928) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713772241
[The Workflow run](https://github.com/apache/airflow/actions/runs/320403262) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713640232
[The Workflow run](https://github.com/apache/airflow/actions/runs/320022045) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-717836498
I thoroughly rewritten the implementation to be compatible with the SSH Operator and also to make the logic much simpler. I think the code is perfect now. I managed to solve all the problems that I encountered and the system tests are passing correctly. Especially SSH tunneling was problematic as SSH tunnel runs in an independent subprocess and I didn't see any error messages so I didn't know you should get authorization separately but now everything works.
I still have a documentation update + some unit tests.
@ephraimbuddy What do you think about the current implementation? Is it clearer?
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713578029
[The Workflow run](https://github.com/apache/airflow/actions/runs/319903176) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713024715
[The Workflow run](https://github.com/apache/airflow/actions/runs/318172808) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713077365
[The Workflow run](https://github.com/apache/airflow/actions/runs/318347559) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713019435
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] stale[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
stale[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-691464741
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713660632
[The Workflow run](https://github.com/apache/airflow/actions/runs/320083422) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713024210
[The Workflow run](https://github.com/apache/airflow/actions/runs/318154240) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713468407
[The Workflow run](https://github.com/apache/airflow/actions/runs/319626549) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713330057
[The Workflow run](https://github.com/apache/airflow/actions/runs/319230004) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713078876
[The Workflow run](https://github.com/apache/airflow/actions/runs/318349535) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713588404
[The Workflow run](https://github.com/apache/airflow/actions/runs/319949826) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713413499
[The Workflow run](https://github.com/apache/airflow/actions/runs/319474887) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713465032
[The Workflow run](https://github.com/apache/airflow/actions/runs/319617519) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713480194
[The Workflow run](https://github.com/apache/airflow/actions/runs/319639661) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on a change in pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r460596997
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,378 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import subprocess
+import time
+from io import StringIO
+from typing import Any, List, Optional, Union
+
+import paramiko
+from google.cloud.oslogin_v1 import OsLoginServiceClient
+from googleapiclient.discovery import build
+from paramiko import AuthenticationException, BadHostKeyException, SSHException
+
+from airflow import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+
+class ComputeEngineSshHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version: The Compute Engine discovery API version
+ :type api_version: str
+ :param gcp_conn_id: The connection id to use when fetching connection information
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+ _conn: Optional[Any] = None
+
+ def __init__(
+ self,
+ api_version: str = 'v1',
+ gcp_conn_id: str = 'google_cloud_default',
+ delegate_to: Optional[str] = None
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+ self._oslogin_conn: Optional[OsLoginServiceClient] = None
+
+ def get_conn(self) -> Any:
+ """
+ Retrieves connection to Google Compute Engine.
+
+ :return: Google Compute Engine services object
+ :rtype: dict
+ """
+ if not self._conn:
+ http_authorized = self._authorize()
+ self._conn = build('compute', self.api_version,
+ http=http_authorized, cache_discovery=False)
+ return self._conn
+
+ def get_oslogin_conn(self) -> OsLoginServiceClient:
+ """
+ Connects to Oslogin API
+ """
+ if not self._oslogin_conn:
+ self._oslogin_conn = OsLoginServiceClient(
+ credentials=self._get_credentials(), client_info=self.client_info
+ )
+
+ return self._oslogin_conn
+
+ def get_sshclient(self,
+ pkey: str,
+ hostname: str,
+ username: str
+ ) -> paramiko.SSHClient:
+ """
+ Get SSHClient
+
+ :param pkey: The private key
+ :type pkey: str
+ :param hostname: The hostname of the target instance
+ :type hostname: str
+ :param username: The username from the login profile of the user account
+ :type username: str
+ """
+
+ client = paramiko.SSHClient()
+ # Default is RejectPolicy
+ # No knownhost checking since we are not storing pkey
+ client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+ client.connect(hostname=hostname,
+ username=username,
+ pkey=pkey,
+ look_for_keys=False)
+ return client
+
+ def get_service_account(self) -> str:
+ """
+ Get the service account id
+ """
+ account: str = self._get_credentials().service_account_email
+ os_login_client = self.get_oslogin_conn()
+ return os_login_client.user_path(user=account)
+
+ def get_username(self, oslogin=True) -> str:
+ """
+ Get the account username
+
+ :param oslogin: Whether to use OsLogin
+ :type oslogin: bool
+ """
+ account = self._get_credentials().service_account_email
+ if oslogin:
+ profile = self.get_oslogin_conn().get_login_profile(name=self.get_service_account())
+ account = profile.posix_accounts[0]
+ return account.username
+ return account.split("@")[0]
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_connection_type(self, project_id: Optional[str], instance: str, zone: str,
+ hostname: Optional[str] = None,
+ use_iap_tunnel: bool = False,
+ use_internal_ip: bool = False
+ ) -> str:
+ """
+ Selects how to connect to the target instance
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param hostname: The hostname of the target instance. If provided, this would
+ be used, regardless of the setting for use_iap_tunnel and use_internal_ip
+ :type hostname: str
+ :param use_iap_tunnel: Whether to connect through IAP tunnel
+ :type use_iap_tunnel: bool
+ :param use_internal_ip: Whether to connect using internal IP
+ :type use_internal_ip: bool
+ """
+ if not hostname:
+ instance_info = self.instance_info(project_id=project_id,
+ zone=zone,
+ instance=instance)
+ if use_iap_tunnel or use_internal_ip:
+ return instance_info["networkInterfaces"][0] \
+ .get("networkIP")
+ else:
+ access_config = instance_info \
+ .get("networkInterfaces")[0] \
+ .get("accessConfigs")
+ if access_config:
+ return access_config[0].get("natIP")
+ else:
+ raise AirflowException("The target instance does not have external IP,"
+ "consider specifying use_internal_ip or use_iap_tunnel")
+ return hostname
+
+ def execute(self, cmd: Union[str, List[str]]) -> None:
+ """
+ Executes command in local machine
+
+ :param cmd: The command to execute
+ :type cmd: Union[str, List[str]]
+ """
+ self.log.info('Executing command: %s', str(cmd))
+ process = subprocess.Popen(cmd)
+ output = process.communicate()[0]
+ returncode = process.returncode
+ if returncode:
+ # Error
+ self.log.error('Command returned error status %s', returncode)
+ if output:
+ self.log.info(output)
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def instance_info(self, project_id, zone, instance):
+ """
+ Used to get instance information
+
+ :param project_id: The project ID where the instance is located
+ :type project_id: str
+ :param zone: The zone of the instance
+ :type zone: str
+ :param instance: The name of the instance
+ :type instance: str
+ """
+ instance_info = self.get_conn().instances().get( # pylint: disable=no-member
+ project=project_id,
+ instance=instance,
+ zone=zone).execute()
+ return instance_info
+
+ def _update_metadata(self, keys, metadata):
+ """
+ Updates ssh keys on metadata dictionary
+
+ :param keys: The public key to set on the metadata
+ :type keys: str
+ :param metadata: A dictionary containing instance metadata
+ :type metadata: dict
+ """
+ items = metadata.get("items")
+ for item in items:
+ if item.get("key") == "ssh-keys":
+ keys += item["value"]
+ item['value'] = keys
+ break
+ else:
+ new_dict = dict(key='ssh-keys', value=keys)
+ metadata['items'] = [new_dict]
+ return metadata
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def set_sshkey_metadata(self, project_id, zone, instance, username, pubkey):
+ """
+ Sets ssh key on instance metadata
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param username: The logged in account username
+ :type username: str
+ :param pubkey: The public key to set on the instance
+ :type pubkey: str
+ """
+ instance_info = self.instance_info(project_id=project_id,
+ zone=zone,
+ instance=instance)
+ keys = username + ":" + pubkey + "\n"
+ metadata = instance_info['metadata']
+ metadata = self._update_metadata(keys, metadata)
+ self.get_conn().instances().setMetadata( # pylint: disable=no-member
+ project=project_id,
+ zone=zone,
+ instance=instance,
+ body=metadata
+ ).execute()
+ time.sleep(10) # TODO improve on this. The status should be Done before proceeding
+
+ def set_oslogin_key(self, pubkey, expire_time):
+ """
+ Sets public key on instance os login
+
+ :param pubkey: The public key to set on the OsLogin
+ :type pubkey: str
+ :param expire_time: The amount of time before the key expires in seconds
+ :type expire_time: int
+ """
+ expiration = int((time.time() + expire_time) * 1000000)
+ body = {
+ "key": pubkey,
+ "expiration_time_usec": expiration
+ }
+ os_login_client = self.get_oslogin_conn()
+ os_login_client.import_ssh_public_key(
+ parent=self.get_service_account(),
+ ssh_public_key=body
+ )
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def _create_ssh_key(self,
+ instance: str,
+ zone: str,
+ project_id: Optional[str] = None,
+ expire_time: int = 300,
+ use_oslogin: bool = True) -> Any:
+ """
+ Generate a temporary SSH key and apply it to the specified account
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param expire_time: The maximum amount of time before the private key is expired
+ :type expire_time: int
+ :param use_oslogin: Whether to use OsLogin to manage key
+ :type use_oslogin: int
+ """
+ username = self.get_username(oslogin=use_oslogin)
+ try:
+ self.log.info("Generating ssh keys...")
+ pkey_file = StringIO()
+ pkey_obj = paramiko.RSAKey.generate(2048)
+ pkey_obj.write_private_key(pkey_file)
+ pubkey = f"{pkey_obj.get_name()} {pkey_obj.get_base64()} {username}"
+ except (IOError, SSHException) as err:
+ raise AirflowException(f"Error encountered creating ssh keys, {err}")
+
+ if use_oslogin:
+ # Expiration time is in microseconds.
+ self.set_oslogin_key(pubkey, expire_time)
+ else:
+ # using instance metadata to manage keys
+ self.set_sshkey_metadata(project_id=project_id,
+ zone=zone,
+ instance=instance,
+ username=username,
+ pubkey=pubkey)
+ return pkey_obj
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def run(self,
+ cmd: str, instance: str, zone: str,
+ project_id: Optional[str] = None,
+ expire_time: int = 300,
+ hostname: Optional[str] = None,
+ use_internal_ip: bool = False,
+ use_iap_tunnel: bool = False,
+ use_oslogin: bool = True
+ ) -> Any:
+ """
+ Execute a command in a remote instance
+
+ :param cmd: The command to execute on the instance
+ :type cmd: str
+ :param instance: The name of the instance the command will be executed on
+ :type instance: str
+ :param zone: The zone where the instance is located
+ :type zone: str
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param expire_time: The maximum amount of time before the private key is expired
+ :type expire_time: int
+ :param hostname: The hostname of the instance.
+ :type hostname: str
+ :param use_iap_tunnel: Whether to connect through IAP tunnel
+ :type use_iap_tunnel: bool
+ :param use_internal_ip: Whether to connect using internal IP
+ :type use_internal_ip: bool
+ :param use_oslogin: Whether to manage keys using OsLogin API. If false,
+ keys are managed using instance metadata
+ """
+ hostname = self.get_connection_type(project_id=project_id,
+ instance=instance,
+ zone=zone,
+ hostname=hostname,
+ use_iap_tunnel=use_iap_tunnel,
+ use_internal_ip=use_internal_ip
+ )
+
+ username = self.get_username(oslogin=use_oslogin)
+
+ self.log.info("Creating remote connection to host")
+ pkey = self._create_ssh_key(
+ instance=instance,
+ project_id=project_id,
+ zone=zone,
+ expire_time=expire_time,
+ use_oslogin=use_oslogin
+ )
+ sshclient = self.get_sshclient(pkey=pkey,
Review comment:
There is one problem if internal IP addresses are used. Two servers can use the same IP address if they are on two different networks. We should make sure that we have connected to the expected server. You can do this by running a simple script on a remote server.
```
metadata_id_url = 'http://metadata.google.internal/computeMetadata/v1/instance/id'
remote_command = f'[ `curl "{metadata_id_url}" -H "Metadata-Flavor: Google" -q` = {instance_id} ] || exit 42'
```
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713177999
[The Workflow run](https://github.com/apache/airflow/actions/runs/318665882) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713026949
[The Workflow run](https://github.com/apache/airflow/actions/runs/318189202) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713196836
[The Workflow run](https://github.com/apache/airflow/actions/runs/318737865) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] turbaszek commented on a change in pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
turbaszek commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r516570816
##########
File path: airflow/providers/google/cloud/hooks/compute.py
##########
@@ -378,3 +378,79 @@ def _check_global_operation_status(
.get(project=project_id, operation=operation_name)
.execute(num_retries=num_retries)
)
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_instance_info(self, zone: str, resource_id: str, project_id: str):
Review comment:
```suggestion
def get_instance_info(self, zone: str, resource_id: str, project_id: str) -> Dict[str, Any]:
```
##########
File path: airflow/providers/google/cloud/hooks/compute.py
##########
@@ -378,3 +378,79 @@ def _check_global_operation_status(
.get(project=project_id, operation=operation_name)
.execute(num_retries=num_retries)
)
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_instance_info(self, zone: str, resource_id: str, project_id: str):
+ """
+ Gets instance information.
+
+ :param zone: Google Cloud zone where the Instance Group Manager exists
+ :type zone: str
+ :param resource_id: Name of the Instance Group Manager
+ :type resource_id: str
+ :param project_id: Optional, Google Cloud project ID where the
+ Compute Engine Instance exists. If set to None or missing,
+ the default project_id from the Google Cloud connection is used.
+ :type project_id: str
+ """
+ instance_info = (
+ self.get_conn() # pylint: disable=no-member
+ .instances()
+ .get(project=project_id, instance=resource_id, zone=zone)
+ .execute(num_retries=self.num_retries)
+ )
+ return instance_info
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_instance_address(
+ self, zone: str, resource_id: str, project_id: str, use_internal_ip: bool = False
+ ) -> str:
+ """
+ Return network address associated to instance.
+
+ :param zone: Google Cloud zone where the Instance Group Manager exists
+ :type zone: str
+ :param resource_id: Name of the Instance Group Manager
+ :type resource_id: str
+ :param project_id: Optional, Google Cloud project ID where the
+ Compute Engine Instance exists. If set to None or missing,
+ the default project_id from the Google Cloud connection is used.
+ :type project_id: str
+ :param use_internal_ip: If true, return private IP address.
+ :type use_internal_ip: bool
+ """
+ instance_info = self.get_instance_info(project_id=project_id, resource_id=resource_id, zone=zone)
+ if use_internal_ip:
+ return instance_info["networkInterfaces"][0].get("networkIP")
+
+ access_config = instance_info.get("networkInterfaces")[0].get("accessConfigs")
+ if access_config:
+ return access_config[0].get("natIP")
+ raise AirflowException("The target instance does not have external IP")
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def set_instance_metadata(self, zone: str, resource_id: str, metadata: Dict[str, str], project_id: str):
Review comment:
```suggestion
def set_instance_metadata(self, zone: str, resource_id: str, metadata: Dict[str, str], project_id: str) -> None:
```
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,304 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import shlex
+import time
+from io import StringIO
+from typing import Any, Optional
+
+import paramiko
+from cached_property import cached_property
+from google.api_core.retry import exponential_sleep_generator
+
+from airflow import AirflowException
+from airflow.providers.google.cloud.hooks.compute import ComputeEngineHook
+from airflow.providers.google.cloud.hooks.os_login import OSLoginHook
+from airflow.providers.ssh.hooks.ssh import SSHHook
+
+
+class _GCloudAuthorizedSSHClient(paramiko.SSHClient):
+ """SSH Client that maintains the context for gcloud authorization during the connection"""
+
+ def __init__(self, google_hook, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.ssh_client = paramiko.SSHClient()
+ self.google_hook = google_hook
+ self.decorator = None
+
+ def connect(self, *args, **kwargs): # pylint: disable=signature-differs
+ self.decorator = self.google_hook.provide_authorized_gcloud()
+ self.decorator.__enter__()
+ return super().connect(*args, **kwargs)
+
+ def close(self):
+ if self.decorator:
+ self.decorator.__exit__(None, None, None)
+ self.decorator = None
+ return super().close()
+
+ def __exit__(self, type_, value, traceback):
+ if self.decorator:
+ self.decorator.__exit__(type_, value, traceback)
+ self.decorator = None
+ return super().__exit__(type_, value, traceback)
+
+
+class ComputeEngineSSHHook(SSHHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param instance_name: The name of the Compute Engine instance
+ :type instance_name: str
+ :param zone: The zone of the Compute Engine instance
+ :type zone: str
+ :param user: The name of the user on which the login attempt will be made
+ :type user: str
+ :param project_id: The project ID of the remote instance
+ :type project_id: str
+ :param gcp_conn_id: The connection id to use when fetching connection info
+ :type gcp_conn_id: str
+ :param hostname: The hostname of the target instance. If it is not passed, it will be detected
+ automatically.
+ :type hostname: str
+ :param use_iap_tunnel: Whether to connect through IAP tunnel
+ :type use_iap_tunnel: bool
+ :param use_internal_ip: Whether to connect using internal IP
+ :type use_internal_ip: bool
+ :param use_oslogin: Whether to manage keys using OsLogin API. If false,
+ keys are managed using instance metadata
+ :param expire_time: The maximum amount of time in seconds before the private key expires
+ :type expire_time: int
+ :param gcp_conn_id: The connection id to use when fetching connection information
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+
+ def __init__( # pylint: disable=too-many-arguments
+ self,
+ gcp_conn_id: str = 'google_cloud_default',
+ instance_name: Optional[str] = None,
+ zone: Optional[str] = None,
+ user: Optional[str] = 'root',
+ project_id: Optional[str] = None,
+ hostname: Optional[str] = None,
+ use_internal_ip: bool = False,
+ use_iap_tunnel: bool = False,
+ use_oslogin: bool = True,
+ expire_time: int = 300,
+ delegate_to: Optional[str] = None,
+ ) -> None:
+ # Ignore original constructor
+ # super().__init__() # pylint: disable=super-init-not-called
+ self.instance_name = instance_name
+ self.zone = zone
+ self.user = user
+ self.project_id = project_id
+ self.hostname = hostname
+ self.use_internal_ip = use_internal_ip
+ self.use_iap_tunnel = use_iap_tunnel
+ self.use_oslogin = use_oslogin
+ self.expire_time = expire_time
+ self.gcp_conn_id = gcp_conn_id
+ self.delegate_to = delegate_to
+ self._conn: Optional[Any] = None
+
+ @cached_property
+ def _oslogin_hook(self) -> OSLoginHook:
+ return OSLoginHook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to)
+
+ @cached_property
+ def _compute_hook(self) -> ComputeEngineHook:
+ return ComputeEngineHook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to)
+
+ def _load_connection_config(self):
+ def _boolify(value):
+ if isinstance(value, bool):
+ return value
+ if isinstance(value, str):
+ if value.lower() == 'false':
+ return False
+ elif value.lower() == 'true':
+ return True
+ return False
+
+ def intify(key, value, default):
+ if value is None:
+ return default
+ if isinstance(value, str) and value.strip() == '':
+ return default
+ try:
+ return int(value)
+ except ValueError:
+ raise AirflowException(
+ f"The {key} field should be a integer. "
+ f"Current value: \"{value}\" (type: {type(value)}). "
+ f"Please check the connection configuration."
+ )
+
+ conn = self.get_connection(self.gcp_conn_id)
+ if conn and conn.conn_type == "gcpssh":
Review comment:
Should this connection be configurable?
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,304 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import shlex
+import time
+from io import StringIO
+from typing import Any, Optional
+
+import paramiko
+from cached_property import cached_property
+from google.api_core.retry import exponential_sleep_generator
+
+from airflow import AirflowException
+from airflow.providers.google.cloud.hooks.compute import ComputeEngineHook
+from airflow.providers.google.cloud.hooks.os_login import OSLoginHook
+from airflow.providers.ssh.hooks.ssh import SSHHook
+
+
+class _GCloudAuthorizedSSHClient(paramiko.SSHClient):
+ """SSH Client that maintains the context for gcloud authorization during the connection"""
+
+ def __init__(self, google_hook, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.ssh_client = paramiko.SSHClient()
+ self.google_hook = google_hook
+ self.decorator = None
+
+ def connect(self, *args, **kwargs): # pylint: disable=signature-differs
+ self.decorator = self.google_hook.provide_authorized_gcloud()
+ self.decorator.__enter__()
+ return super().connect(*args, **kwargs)
+
+ def close(self):
+ if self.decorator:
+ self.decorator.__exit__(None, None, None)
+ self.decorator = None
+ return super().close()
+
+ def __exit__(self, type_, value, traceback):
+ if self.decorator:
+ self.decorator.__exit__(type_, value, traceback)
+ self.decorator = None
+ return super().__exit__(type_, value, traceback)
+
+
+class ComputeEngineSSHHook(SSHHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param instance_name: The name of the Compute Engine instance
+ :type instance_name: str
+ :param zone: The zone of the Compute Engine instance
+ :type zone: str
+ :param user: The name of the user on which the login attempt will be made
+ :type user: str
+ :param project_id: The project ID of the remote instance
+ :type project_id: str
+ :param gcp_conn_id: The connection id to use when fetching connection info
+ :type gcp_conn_id: str
+ :param hostname: The hostname of the target instance. If it is not passed, it will be detected
+ automatically.
+ :type hostname: str
+ :param use_iap_tunnel: Whether to connect through IAP tunnel
+ :type use_iap_tunnel: bool
+ :param use_internal_ip: Whether to connect using internal IP
+ :type use_internal_ip: bool
+ :param use_oslogin: Whether to manage keys using OsLogin API. If false,
+ keys are managed using instance metadata
+ :param expire_time: The maximum amount of time in seconds before the private key expires
+ :type expire_time: int
+ :param gcp_conn_id: The connection id to use when fetching connection information
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+
+ def __init__( # pylint: disable=too-many-arguments
+ self,
+ gcp_conn_id: str = 'google_cloud_default',
+ instance_name: Optional[str] = None,
+ zone: Optional[str] = None,
+ user: Optional[str] = 'root',
+ project_id: Optional[str] = None,
+ hostname: Optional[str] = None,
+ use_internal_ip: bool = False,
+ use_iap_tunnel: bool = False,
+ use_oslogin: bool = True,
+ expire_time: int = 300,
+ delegate_to: Optional[str] = None,
+ ) -> None:
+ # Ignore original constructor
+ # super().__init__() # pylint: disable=super-init-not-called
+ self.instance_name = instance_name
+ self.zone = zone
+ self.user = user
+ self.project_id = project_id
+ self.hostname = hostname
+ self.use_internal_ip = use_internal_ip
+ self.use_iap_tunnel = use_iap_tunnel
+ self.use_oslogin = use_oslogin
+ self.expire_time = expire_time
+ self.gcp_conn_id = gcp_conn_id
+ self.delegate_to = delegate_to
+ self._conn: Optional[Any] = None
+
+ @cached_property
+ def _oslogin_hook(self) -> OSLoginHook:
+ return OSLoginHook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to)
+
+ @cached_property
+ def _compute_hook(self) -> ComputeEngineHook:
+ return ComputeEngineHook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to)
+
+ def _load_connection_config(self):
+ def _boolify(value):
+ if isinstance(value, bool):
+ return value
+ if isinstance(value, str):
+ if value.lower() == 'false':
+ return False
+ elif value.lower() == 'true':
+ return True
+ return False
+
+ def intify(key, value, default):
+ if value is None:
+ return default
+ if isinstance(value, str) and value.strip() == '':
+ return default
+ try:
+ return int(value)
+ except ValueError:
+ raise AirflowException(
+ f"The {key} field should be a integer. "
+ f"Current value: \"{value}\" (type: {type(value)}). "
+ f"Please check the connection configuration."
+ )
+
+ conn = self.get_connection(self.gcp_conn_id)
+ if conn and conn.conn_type == "gcpssh":
+ self.instance_name = self._compute_hook._get_field( # pylint: disable=protected-access
+ "instance_name", self.instance_name
+ )
+ self.zone = self._compute_hook._get_field("zone", self.zone) # pylint: disable=protected-access
+ self.user = conn.login if conn.login else self.user
+ # self.project_id is skipped intentionally
+ self.hostname = conn.host if conn.host else self.hostname
+ self.use_internal_ip = _boolify(
+ self._compute_hook._get_field("use_internal_ip") # pylint: disable=protected-access
+ )
+ self.use_iap_tunnel = _boolify(
+ self._compute_hook._get_field("use_iap_tunnel") # pylint: disable=protected-access
+ )
+ self.use_oslogin = _boolify(
+ self._compute_hook._get_field("use_oslogin") # pylint: disable=protected-access
+ )
+ self.expire_time = intify(
+ "expire_time",
+ self._compute_hook._get_field("expire_time"), # pylint: disable=protected-access
+ self.expire_time,
+ )
+
+ def get_conn(self) -> paramiko.SSHClient:
+ """Return SSH connection."""
+ self._load_connection_config()
+ if not self.project_id:
+ self.project_id = self._compute_hook.project_id
+
+ missing_fields = [k for k in ["instance_name", "zone", "project_id"] if not getattr(self, k)]
+ if not self.instance_name or not self.zone or not self.project_id:
+ raise AirflowException(
+ f"Required parameters are missing: {missing_fields}. These parameters be passed either as "
+ "keyword parameter or as extra field in Airfow connection definition. Both are not set!"
+ )
+
+ self.log.info(
+ "Connecting to instance: instance_name=%s, user=%s, zone=%s, "
+ "use_internal_ip=%s, use_iap_tunnel=%s, use_os_login=%s",
+ self.instance_name,
+ self.user,
+ self.zone,
+ self.use_internal_ip,
+ self.use_iap_tunnel,
+ self.use_oslogin,
+ )
+ if not self.hostname:
+ hostname = self._compute_hook.get_instance_address(
+ zone=self.zone,
+ resource_id=self.instance_name,
+ project_id=self.project_id,
+ use_internal_ip=self.use_internal_ip or self.use_iap_tunnel,
+ )
+ else:
+ hostname = self.hostname
+
+ privkey, pubkey = self._generate_ssh_key(self.user)
+ if self.use_oslogin:
+ user = self._authorize_os_login(pubkey)
+ else:
+ user = self.user
+ self._authorize_compute_engine_instance_metadata(pubkey)
+
+ proxy_command = None
+ if self.use_iap_tunnel:
+ proxy_command_args = [
+ 'gcloud',
+ 'compute',
+ 'start-iap-tunnel',
+ str(self.instance_name),
+ '22',
+ '--listen-on-stdin',
+ f'--project={self.project_id}',
+ f'--zone={self.zone}',
+ '--verbosity=warning',
+ ]
+ proxy_command = " ".join(shlex.quote(arg) for arg in proxy_command_args)
+
+ sshclient = self._connect_to_instance(user, hostname, privkey, proxy_command)
+ return sshclient
+
+ def _connect_to_instance(self, user, hostname, pkey, proxy_command) -> paramiko.SSHClient:
+ self.log.info("Opening remote connection to host: username=%s, hostname=%s", user, hostname)
+ for time_to_wait in exponential_sleep_generator(initial=1, maximum=10):
+ try:
+ client = _GCloudAuthorizedSSHClient(self._compute_hook)
+ # Default is RejectPolicy
+ # No known host checking since we are not storing privatekey
+ client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+ client.connect(
+ hostname=hostname,
+ username=user,
+ pkey=pkey,
+ sock=paramiko.ProxyCommand(proxy_command) if proxy_command else None,
+ look_for_keys=False,
+ )
+ return client
+ except paramiko.SSHException:
+ if time_to_wait == 10:
Review comment:
Can you explain it or leave a comment?
##########
File path: airflow/providers/ssh/operators/ssh.py
##########
@@ -89,7 +90,8 @@ def execute(self, context) -> Union[bytes, str, bool]:
self.log.info(
"ssh_hook is not provided or invalid. " "Trying ssh_conn_id to create SSHHook."
)
- self.ssh_hook = SSHHook(ssh_conn_id=self.ssh_conn_id, timeout=self.timeout)
+ self.ssh_hook = BaseHook.get_connection(self.ssh_conn_id)
Review comment:
Is the `ssh_hook` proper name? I think this will be a connection now?
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713553451
[The Workflow run](https://github.com/apache/airflow/actions/runs/319848075) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-712995978
[The Workflow run](https://github.com/apache/airflow/actions/runs/317999785) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713013595
[The Workflow run](https://github.com/apache/airflow/actions/runs/318053201) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] stale[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
stale[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-691464741
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713778757
[The Workflow run](https://github.com/apache/airflow/actions/runs/320427374) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713223096
[The Workflow run](https://github.com/apache/airflow/actions/runs/318839594) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713381436
[The Workflow run](https://github.com/apache/airflow/actions/runs/319384708) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713123652
[The Workflow run](https://github.com/apache/airflow/actions/runs/318484726) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] stale[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
stale[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-691464741
This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713061395
[The Workflow run](https://github.com/apache/airflow/actions/runs/318297415) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713724307
[The Workflow run](https://github.com/apache/airflow/actions/runs/320207125) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713495177
[The Workflow run](https://github.com/apache/airflow/actions/runs/319717477) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-712868289
[The Workflow run](https://github.com/apache/airflow/actions/runs/317669593) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713019076
[The Workflow run](https://github.com/apache/airflow/actions/runs/318086772) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] ephraimbuddy commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
ephraimbuddy commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-664800507
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-712891541
[The Workflow run](https://github.com/apache/airflow/actions/runs/317751136) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713642253
[The Workflow run](https://github.com/apache/airflow/actions/runs/320033364) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on a change in pull request #9879: [WIP]Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r458446010
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,250 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import subprocess
+import time
+import uuid
+from typing import Any, List, Optional, Union
+
+import paramiko
+from googleapiclient.discovery import build
+from paramiko import SSHException
+
+from airflow import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+
+class ComputeEngineSshHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version: The Oslogin discovery API version
+ :type api_version: str
+ :param gcp_conn_id: The connection id to use when fetching connection info
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+
+ def __init__(self,
+ api_version: str = "v1",
+ gcp_conn_id: str = "google_cloud_default",
+ delegate_to: Optional[str] = None
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+ self.conn = None
+
+ def get_conn(self) -> Any:
+ """
+ Connects to Oslogin API
+ """
+ if self.conn is None:
+ http_authorized = self._authorize()
+ self.conn = build("oslogin", self.api_version,
+ http=http_authorized, cache_discovery=False)
+ return self.conn
+
+ def get_sshclient(self,
+ private_key_file: str,
+ hostname: str,
+ username: str,
+ strict_host_checking: bool = True
+ ) -> paramiko.SSHClient:
+ """
+ Get SSHClient
+
+ :param private_key_file: The filename of the private key
+ :type private_key_file: str
+ :param hostname: The hostname of the target instance
+ :type hostname: str
+ :param username: The POXIS username from the login profile of the user account
+ :type username: str
+ :param strict_host_checking: Whether to verify key
+ :type strict_host_checking: bool
+ """
+
+ self.log.info("Creating remote connection to host")
+ client = paramiko.SSHClient()
+ if not strict_host_checking:
+ self.log.warning('No Host Key Verification. This wont protect '
+ 'against Man-In-The-Middle attacks')
+ # Default is RejectPolicy
+ client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+ client.connect(hostname=hostname,
+ username=username,
+ key_filename=private_key_file)
+ return client
+
+ def get_service_account(self) -> str:
+ """
+ Get the service account id
+ """
+ account: str = self._get_credentials().service_account_email
+ if not account.startswith('users/'):
+ account = 'users/' + account
+ return account
+
+ def get_username(self) -> str:
+ """
+ Using the OS Login API, get the POSIX username from the login profile
+ for the service account.
+
+ """
+ conn = self.get_conn()
+ account = self.get_service_account()
+ profile = conn.users().getLoginProfile(name=account).execute_in_compute() # pylint: disable=no-member
+ username: str = profile.get('posixAccounts')[0].get('username')
+ return username
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_hostname(self, project_id: Optional[str], instance: str, zone: str,
+ hostname: Optional[str] = None) -> str:
+ """
+ Create the target instance hostname
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param hostname: The hostname of the target instance. If not provided,
+ it is generated from zone, instance name and project id
+ :type hostname: str
+ """
+ if not hostname:
+ hostname = f'{instance}.{zone}.c.{project_id}.internal'
Review comment:
Hostname is a bit complicated in this case, but we definitely shouldn't use a domain name, but IP addresses.
If user passes hostname, we should use the value provided by user
If user turns on IAP tunnel, we should use a private IP address.
If user enforce private IP (`use_internal_ip` option`, then we should use private IP
Otherwise, we should use the public IP address.
Public IP address is a first non-empty value in `networkInterfaces[].accessConfigs[].natIP` attribute of [Instance object](https://cloud.google.com/compute/docs/reference/rest/v1/instances/get).
Private IP address is a first non-empty value in `networkInterfaces[].networkIP` attribute of [Instance object](https://cloud.google.com/compute/docs/reference/rest/v1/instances/get).
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on a change in pull request #9879: [WIP]Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r459459626
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,327 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import pprint
+import subprocess
+import time
+import uuid
+from io import StringIO
+from typing import Any, List, Optional, Tuple, Union
+
+import paramiko
+from google.cloud.oslogin_v1 import OsLoginServiceClient
+from googleapiclient.discovery import build
+from paramiko import SSHException
+
+from airflow import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+
+class ComputeEngineSshHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version: The Compute Engine discovery API version
+ :type api_version: str
+ :param gcp_conn_id: The connection id to use when fetching connection info
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+ _oslogin_conn: Optional[OsLoginServiceClient] = None
+ _sshclient: Optional[paramiko.SSHClient] = None
+ _conn: Optional[Any] = None
+
+ def __init__(
+ self,
+ api_version: str = 'v1',
+ gcp_conn_id: str = 'google_cloud_default',
+ delegate_to: Optional[str] = None
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+
+ def get_conn(self) -> Any:
+ """
+ Retrieves connection to Google Compute Engine.
+
+ :return: Google Compute Engine services object
+ :rtype: dict
+ """
+ if not self._conn:
+ http_authorized = self._authorize()
+ self._conn = build('compute', self.api_version,
+ http=http_authorized, cache_discovery=False)
+ return self._conn
+
+ def get_oslogin_conn(self) -> OsLoginServiceClient:
+ """
+ Connects to Oslogin API
+ """
+ if not self._oslogin_conn:
+ self._oslogin_conn = OsLoginServiceClient(
+ credentials=self._get_credentials(), client_info=self.client_info
+ )
+
+ return self._oslogin_conn
+
+ def get_sshclient(self,
+ private_key_file: str,
+ hostname: str,
+ username: str
+ ) -> paramiko.SSHClient:
+ """
+ Get SSHClient
+
+ :param private_key_file: The private key filename
+ :type private_key_file: str
+ :param hostname: The hostname of the target instance
+ :type hostname: str
+ :param username: The POXIS username from the login profile of the user account
+ :type username: str
+ """
+
+ if not self._sshclient:
+ self._sshclient = paramiko.SSHClient()
+ # Default is RejectPolicy
+ # No knownhost checking since we are not storing pkey
+ self._sshclient.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+ self._sshclient.connect(hostname=hostname,
+ username=username,
+ key_filename=private_key_file,
+ look_for_keys=False)
+ return self._sshclient
+
+ def get_service_account(self) -> str:
+ """
+ Get the service account id
+ """
+ account: str = self._get_credentials().service_account_email
+ os_login_client = self.get_oslogin_conn()
+ return os_login_client.user_path(user=account)
+
+ def get_username(self, oslogin=True) -> str:
+ """
+ Get the account username
+ """
+ account = self._get_credentials().service_account_email
+ if oslogin:
+ profile = self.get_oslogin_conn().get_login_profile(name=self.get_service_account())
+ account = profile.posix_accounts[0]
+ return account.username
+ return account.split("@")[0]
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_connection_type(self, project_id: Optional[str], instance: str, zone: str,
+ hostname: Optional[str] = None,
+ use_iap_tunnel: bool = False,
+ use_internal_ip: bool = False
+ ) -> str:
+ """
+ Selects how to connect to the target instance
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param hostname: The hostname of the target instance. If provided, this would
+ be used, regardless of the setting for use_iap_tunnel and use_internal_ip
+ :type hostname: str
+ :param use_iap_tunnel: Whether to connect through IAP tunnel
+ :type use_iap_tunnel: bool
+ :param use_internal_ip: Whether to connect using internal IP
+ :type use_internal_ip: bool
+ """
+ if not hostname:
+ instance_info = self.get_conn() \
+ .instances().get(project=project_id,
+ zone=zone,
+ instance=instance).execute()
+ if use_iap_tunnel or use_internal_ip:
+ return instance_info["networkInterfaces"][0] \
+ .get("networkIP")
+ else:
+ external_ip = instance_info \
+ .get("networkInterfaces")[0] \
+ .get("accessConfigs")
+ if external_ip:
+ return external_ip[0].get("natIP")
Review comment:
When natIP is empty, we should look at the access configuration.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713194865
[The Workflow run](https://github.com/apache/airflow/actions/runs/318729886) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-712912334
[The Workflow run](https://github.com/apache/airflow/actions/runs/317881425) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on a change in pull request #9879: [WIP]Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r458455540
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,250 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import subprocess
+import time
+import uuid
+from typing import Any, List, Optional, Union
+
+import paramiko
+from googleapiclient.discovery import build
+from paramiko import SSHException
+
+from airflow import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+
+class ComputeEngineSshHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version: The Oslogin discovery API version
+ :type api_version: str
+ :param gcp_conn_id: The connection id to use when fetching connection info
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+
+ def __init__(self,
+ api_version: str = "v1",
+ gcp_conn_id: str = "google_cloud_default",
+ delegate_to: Optional[str] = None
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+ self.conn = None
+
+ def get_conn(self) -> Any:
+ """
+ Connects to Oslogin API
+ """
+ if self.conn is None:
+ http_authorized = self._authorize()
+ self.conn = build("oslogin", self.api_version,
+ http=http_authorized, cache_discovery=False)
+ return self.conn
+
+ def get_sshclient(self,
+ private_key_file: str,
+ hostname: str,
+ username: str,
+ strict_host_checking: bool = True
+ ) -> paramiko.SSHClient:
+ """
+ Get SSHClient
+
+ :param private_key_file: The filename of the private key
+ :type private_key_file: str
+ :param hostname: The hostname of the target instance
+ :type hostname: str
+ :param username: The POXIS username from the login profile of the user account
+ :type username: str
+ :param strict_host_checking: Whether to verify key
+ :type strict_host_checking: bool
+ """
+
+ self.log.info("Creating remote connection to host")
+ client = paramiko.SSHClient()
+ if not strict_host_checking:
+ self.log.warning('No Host Key Verification. This wont protect '
+ 'against Man-In-The-Middle attacks')
+ # Default is RejectPolicy
+ client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+ client.connect(hostname=hostname,
+ username=username,
+ key_filename=private_key_file)
+ return client
+
+ def get_service_account(self) -> str:
+ """
+ Get the service account id
+ """
+ account: str = self._get_credentials().service_account_email
+ if not account.startswith('users/'):
+ account = 'users/' + account
+ return account
+
+ def get_username(self) -> str:
+ """
+ Using the OS Login API, get the POSIX username from the login profile
+ for the service account.
+
+ """
+ conn = self.get_conn()
+ account = self.get_service_account()
+ profile = conn.users().getLoginProfile(name=account).execute_in_compute() # pylint: disable=no-member
+ username: str = profile.get('posixAccounts')[0].get('username')
+ return username
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_hostname(self, project_id: Optional[str], instance: str, zone: str,
+ hostname: Optional[str] = None) -> str:
+ """
+ Create the target instance hostname
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param hostname: The hostname of the target instance. If not provided,
+ it is generated from zone, instance name and project id
+ :type hostname: str
+ """
+ if not hostname:
+ hostname = f'{instance}.{zone}.c.{project_id}.internal'
+ return hostname
+
+ def execute_local(self, cmd: Union[str, List[str]]) -> None:
+ """
+ Executes command in local machine
+
+ :param cmd: The command to execute
+ :type cmd: Union[str, List[str]]
+ """
+ self.log.info('Executing command: %s', str(cmd))
+ process = subprocess.Popen(cmd)
+ output = process.communicate()[0]
+ return_code = process.returncode
+ if return_code:
+ # Error
+ raise AirflowException(f'The Command {cmd} returned error status {return_code}')
+ if output:
+ self.log.info(output)
+
+ def create_ssh_key(self,
+ account: str,
+ expire_time: int = 300) -> str:
+ """
+ Generate a temporary SSH key and apply it to the specified account
+
+ :param account: The service account id
+ :type account: str
+ :param expire_time: The maximum amount of time before the private key is expired
+ :type expire_time: int
+ """
+ try:
+ self.log.info("Generating ssh keys...")
+ private_key_file = '/tmp/key-' + str(uuid.uuid4())
Review comment:
You can also read key from .key attribute and serialize it to string.
https://cryptography.io/en/latest/hazmat/primitives/asymmetric/rsa/#key-serialization
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713533810
[The Workflow run](https://github.com/apache/airflow/actions/runs/319835979) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713445246
[The Workflow run](https://github.com/apache/airflow/actions/runs/319534198) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713622075
[The Workflow run](https://github.com/apache/airflow/actions/runs/320004425) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] ephraimbuddy commented on a change in pull request #9879: [WIP]Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
ephraimbuddy commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r459372463
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,250 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import subprocess
+import time
+import uuid
+from typing import Any, List, Optional, Union
+
+import paramiko
+from googleapiclient.discovery import build
+from paramiko import SSHException
+
+from airflow import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+
+class ComputeEngineSshHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version: The Oslogin discovery API version
+ :type api_version: str
+ :param gcp_conn_id: The connection id to use when fetching connection info
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+
+ def __init__(self,
+ api_version: str = "v1",
+ gcp_conn_id: str = "google_cloud_default",
+ delegate_to: Optional[str] = None
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+ self.conn = None
+
+ def get_conn(self) -> Any:
+ """
+ Connects to Oslogin API
+ """
+ if self.conn is None:
+ http_authorized = self._authorize()
+ self.conn = build("oslogin", self.api_version,
+ http=http_authorized, cache_discovery=False)
+ return self.conn
+
+ def get_sshclient(self,
+ private_key_file: str,
+ hostname: str,
+ username: str,
+ strict_host_checking: bool = True
+ ) -> paramiko.SSHClient:
+ """
+ Get SSHClient
+
+ :param private_key_file: The filename of the private key
+ :type private_key_file: str
+ :param hostname: The hostname of the target instance
+ :type hostname: str
+ :param username: The POXIS username from the login profile of the user account
+ :type username: str
+ :param strict_host_checking: Whether to verify key
+ :type strict_host_checking: bool
+ """
+
+ self.log.info("Creating remote connection to host")
+ client = paramiko.SSHClient()
+ if not strict_host_checking:
+ self.log.warning('No Host Key Verification. This wont protect '
+ 'against Man-In-The-Middle attacks')
+ # Default is RejectPolicy
+ client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+ client.connect(hostname=hostname,
+ username=username,
+ key_filename=private_key_file)
+ return client
+
+ def get_service_account(self) -> str:
+ """
+ Get the service account id
+ """
+ account: str = self._get_credentials().service_account_email
+ if not account.startswith('users/'):
+ account = 'users/' + account
+ return account
+
+ def get_username(self) -> str:
+ """
+ Using the OS Login API, get the POSIX username from the login profile
+ for the service account.
+
+ """
+ conn = self.get_conn()
+ account = self.get_service_account()
+ profile = conn.users().getLoginProfile(name=account).execute_in_compute() # pylint: disable=no-member
+ username: str = profile.get('posixAccounts')[0].get('username')
+ return username
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_hostname(self, project_id: Optional[str], instance: str, zone: str,
+ hostname: Optional[str] = None) -> str:
+ """
+ Create the target instance hostname
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param hostname: The hostname of the target instance. If not provided,
+ it is generated from zone, instance name and project id
+ :type hostname: str
+ """
+ if not hostname:
+ hostname = f'{instance}.{zone}.c.{project_id}.internal'
+ return hostname
+
+ def execute_local(self, cmd: Union[str, List[str]]) -> None:
+ """
+ Executes command in local machine
+
+ :param cmd: The command to execute
+ :type cmd: Union[str, List[str]]
+ """
+ self.log.info('Executing command: %s', str(cmd))
+ process = subprocess.Popen(cmd)
+ output = process.communicate()[0]
+ return_code = process.returncode
+ if return_code:
+ # Error
+ raise AirflowException(f'The Command {cmd} returned error status {return_code}')
+ if output:
+ self.log.info(output)
+
+ def create_ssh_key(self,
+ account: str,
+ expire_time: int = 300) -> str:
+ """
+ Generate a temporary SSH key and apply it to the specified account
+
+ :param account: The service account id
+ :type account: str
+ :param expire_time: The maximum amount of time before the private key is expired
+ :type expire_time: int
+ """
+ try:
+ self.log.info("Generating ssh keys...")
+ private_key_file = '/tmp/key-' + str(uuid.uuid4())
Review comment:
I have tried several ways to use paramiko with StringIO but I always get authentication failure. Using the oslogin and shell command, the system test succeeds but if oslogin is switched off, It doesn't work as expected.
Another issue is the account username, when I check the server, oslogin is using service account unique ID which is a bunch of numbers. Was expecting it to use the part of email before '@' but it uses the unique ID.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713152992
[The Workflow run](https://github.com/apache/airflow/actions/runs/318580734) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713598287
[The Workflow run](https://github.com/apache/airflow/actions/runs/319961905) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713014460
[The Workflow run](https://github.com/apache/airflow/actions/runs/318054171) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713776649
[The Workflow run](https://github.com/apache/airflow/actions/runs/320410803) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713703050
[The Workflow run](https://github.com/apache/airflow/actions/runs/320195050) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-664067949
Have you tested tunneling the ssh connection through the Cloud Identity-Aware Proxy?
From a technical point of view, this means you should use [ProxyCommand](http://docs.paramiko.org/en/stable/api/proxy.html
).
Here is example command:
```bash
gcloud compute start-iap-tunnel \
<INSTANCE_NAME> %p \
--listen-on-stdin \
--project=<PROJECT_ID> \
--zone=<ZONE>
```
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on a change in pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r519460158
##########
File path: airflow/providers/ssh/operators/ssh.py
##########
@@ -89,7 +90,8 @@ def execute(self, context) -> Union[bytes, str, bool]:
self.log.info(
"ssh_hook is not provided or invalid. " "Trying ssh_conn_id to create SSHHook."
)
- self.ssh_hook = SSHHook(ssh_conn_id=self.ssh_conn_id, timeout=self.timeout)
+ self.ssh_hook = BaseHook.get_connection(self.ssh_conn_id)
Review comment:
I have revert this change. I will deal with it in a separate PR.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on a change in pull request #9879: [WIP]Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r456886345
##########
File path: airflow/providers/google/cloud/hooks/compute_engine_ssh.py
##########
@@ -0,0 +1,211 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import subprocess
+import time
+import uuid
+from typing import Any, List, Optional, Union
+
+import requests
+from googleapiclient.discovery import build
+
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+SERVICE_ACCOUNT_METADATA_URL = (
+ 'http://metadata.google.internal/computeMetadata/v1/instance/'
+ 'service-accounts/default/email')
+HEADERS = {'Metadata-Flavor': 'Google'}
+
+
+class ComputeEngineSSHHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version:
+ :param gcp_conn_id:
+ :param delegate_to:
+ :param oslogin:
+ :param account:
+ """
+
+ def __init__(self,
+ api_version: str = "v1",
+ gcp_conn_id: str = "google_cloud_default",
+ delegate_to: Optional[str] = None,
+ oslogin=None,
+ account=None
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+ self.oslogin = oslogin
+ self.account = account
+
+ def get_oslogin(self) -> Any:
+ """
+ Connects to Oslogin API
+ """
+ if self.oslogin is None:
+ http_authorized = self._authorize()
+ self.oslogin = build("oslogin", self.api_version,
+ http=http_authorized, cache_discovery=False)
+ return self.oslogin
+
+ def get_service_account(self) -> str:
+ """
+ Get the service account id
+ :return:
+ """
+ if self.account is None:
+ self.account = requests.get(
+ SERVICE_ACCOUNT_METADATA_URL, headers=HEADERS).text
+ if not self.account.startswith('users/'):
Review comment:
This will only work on GCE instances. Can you try to fetch email from credentials - `hook._get_credentials.service_account_email`? This still won't be perfect as there are edge cases that may not work properly, but we can deal with it later.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj merged pull request #9879: Add Compute Engine SSH hook
Posted by GitBox <gi...@apache.org>.
mik-laj merged pull request #9879:
URL: https://github.com/apache/airflow/pull/9879
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713778957
[The Workflow run](https://github.com/apache/airflow/actions/runs/320430305) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713047290
[The Workflow run](https://github.com/apache/airflow/actions/runs/318245926) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on a change in pull request #9879: [WIP]Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r458454070
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,250 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import subprocess
+import time
+import uuid
+from typing import Any, List, Optional, Union
+
+import paramiko
+from googleapiclient.discovery import build
+from paramiko import SSHException
+
+from airflow import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+
+class ComputeEngineSshHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version: The Oslogin discovery API version
+ :type api_version: str
+ :param gcp_conn_id: The connection id to use when fetching connection info
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+
+ def __init__(self,
+ api_version: str = "v1",
+ gcp_conn_id: str = "google_cloud_default",
+ delegate_to: Optional[str] = None
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+ self.conn = None
+
+ def get_conn(self) -> Any:
+ """
+ Connects to Oslogin API
+ """
+ if self.conn is None:
+ http_authorized = self._authorize()
+ self.conn = build("oslogin", self.api_version,
+ http=http_authorized, cache_discovery=False)
+ return self.conn
+
+ def get_sshclient(self,
+ private_key_file: str,
+ hostname: str,
+ username: str,
+ strict_host_checking: bool = True
+ ) -> paramiko.SSHClient:
+ """
+ Get SSHClient
+
+ :param private_key_file: The filename of the private key
+ :type private_key_file: str
+ :param hostname: The hostname of the target instance
+ :type hostname: str
+ :param username: The POXIS username from the login profile of the user account
+ :type username: str
+ :param strict_host_checking: Whether to verify key
+ :type strict_host_checking: bool
+ """
+
+ self.log.info("Creating remote connection to host")
+ client = paramiko.SSHClient()
+ if not strict_host_checking:
+ self.log.warning('No Host Key Verification. This wont protect '
+ 'against Man-In-The-Middle attacks')
+ # Default is RejectPolicy
+ client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+ client.connect(hostname=hostname,
+ username=username,
+ key_filename=private_key_file)
+ return client
+
+ def get_service_account(self) -> str:
+ """
+ Get the service account id
+ """
+ account: str = self._get_credentials().service_account_email
+ if not account.startswith('users/'):
+ account = 'users/' + account
+ return account
+
+ def get_username(self) -> str:
+ """
+ Using the OS Login API, get the POSIX username from the login profile
+ for the service account.
+
+ """
+ conn = self.get_conn()
+ account = self.get_service_account()
+ profile = conn.users().getLoginProfile(name=account).execute_in_compute() # pylint: disable=no-member
+ username: str = profile.get('posixAccounts')[0].get('username')
+ return username
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_hostname(self, project_id: Optional[str], instance: str, zone: str,
+ hostname: Optional[str] = None) -> str:
+ """
+ Create the target instance hostname
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param hostname: The hostname of the target instance. If not provided,
+ it is generated from zone, instance name and project id
+ :type hostname: str
+ """
+ if not hostname:
+ hostname = f'{instance}.{zone}.c.{project_id}.internal'
+ return hostname
+
+ def execute_local(self, cmd: Union[str, List[str]]) -> None:
+ """
+ Executes command in local machine
+
+ :param cmd: The command to execute
+ :type cmd: Union[str, List[str]]
+ """
+ self.log.info('Executing command: %s', str(cmd))
+ process = subprocess.Popen(cmd)
+ output = process.communicate()[0]
+ return_code = process.returncode
+ if return_code:
+ # Error
+ raise AirflowException(f'The Command {cmd} returned error status {return_code}')
+ if output:
+ self.log.info(output)
+
+ def create_ssh_key(self,
+ account: str,
+ expire_time: int = 300) -> str:
+ """
+ Generate a temporary SSH key and apply it to the specified account
+
+ :param account: The service account id
+ :type account: str
+ :param expire_time: The maximum amount of time before the private key is expired
+ :type expire_time: int
+ """
+ try:
+ self.log.info("Generating ssh keys...")
+ private_key_file = '/tmp/key-' + str(uuid.uuid4())
Review comment:
```python
def generate_key_string(pkey: paramiko.PKey):
key_fh = StringIO()
pkey.write_private_key(key_fh)
key_fh.seek(0)
key_str = key_fh.read()
return key_str
```
You can use this code to prepare SSH key without writing to disk.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713257682
[The Workflow run](https://github.com/apache/airflow/actions/runs/318980139) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-717638440
[The Workflow run](https://github.com/apache/airflow/actions/runs/332636472) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks,^Build docs$,^Spell check docs$,^Backport packages$,^Provider packages,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713643269
[The Workflow run](https://github.com/apache/airflow/actions/runs/320033463) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713434076
[The Workflow run](https://github.com/apache/airflow/actions/runs/319518811) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] ephraimbuddy commented on a change in pull request #9879: [WIP]Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
ephraimbuddy commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r459833486
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,250 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import subprocess
+import time
+import uuid
+from typing import Any, List, Optional, Union
+
+import paramiko
+from googleapiclient.discovery import build
+from paramiko import SSHException
+
+from airflow import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+
+class ComputeEngineSshHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version: The Oslogin discovery API version
+ :type api_version: str
+ :param gcp_conn_id: The connection id to use when fetching connection info
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+
+ def __init__(self,
+ api_version: str = "v1",
+ gcp_conn_id: str = "google_cloud_default",
+ delegate_to: Optional[str] = None
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+ self.conn = None
+
+ def get_conn(self) -> Any:
+ """
+ Connects to Oslogin API
+ """
+ if self.conn is None:
+ http_authorized = self._authorize()
+ self.conn = build("oslogin", self.api_version,
+ http=http_authorized, cache_discovery=False)
+ return self.conn
+
+ def get_sshclient(self,
+ private_key_file: str,
+ hostname: str,
+ username: str,
+ strict_host_checking: bool = True
+ ) -> paramiko.SSHClient:
+ """
+ Get SSHClient
+
+ :param private_key_file: The filename of the private key
+ :type private_key_file: str
+ :param hostname: The hostname of the target instance
+ :type hostname: str
+ :param username: The POXIS username from the login profile of the user account
+ :type username: str
+ :param strict_host_checking: Whether to verify key
+ :type strict_host_checking: bool
+ """
+
+ self.log.info("Creating remote connection to host")
+ client = paramiko.SSHClient()
+ if not strict_host_checking:
+ self.log.warning('No Host Key Verification. This wont protect '
+ 'against Man-In-The-Middle attacks')
+ # Default is RejectPolicy
+ client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+ client.connect(hostname=hostname,
+ username=username,
+ key_filename=private_key_file)
+ return client
+
+ def get_service_account(self) -> str:
+ """
+ Get the service account id
+ """
+ account: str = self._get_credentials().service_account_email
+ if not account.startswith('users/'):
+ account = 'users/' + account
+ return account
+
+ def get_username(self) -> str:
+ """
+ Using the OS Login API, get the POSIX username from the login profile
+ for the service account.
+
+ """
+ conn = self.get_conn()
+ account = self.get_service_account()
+ profile = conn.users().getLoginProfile(name=account).execute_in_compute() # pylint: disable=no-member
+ username: str = profile.get('posixAccounts')[0].get('username')
+ return username
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_hostname(self, project_id: Optional[str], instance: str, zone: str,
+ hostname: Optional[str] = None) -> str:
+ """
+ Create the target instance hostname
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param hostname: The hostname of the target instance. If not provided,
+ it is generated from zone, instance name and project id
+ :type hostname: str
+ """
+ if not hostname:
+ hostname = f'{instance}.{zone}.c.{project_id}.internal'
+ return hostname
+
+ def execute_local(self, cmd: Union[str, List[str]]) -> None:
+ """
+ Executes command in local machine
+
+ :param cmd: The command to execute
+ :type cmd: Union[str, List[str]]
+ """
+ self.log.info('Executing command: %s', str(cmd))
+ process = subprocess.Popen(cmd)
+ output = process.communicate()[0]
+ return_code = process.returncode
+ if return_code:
+ # Error
+ raise AirflowException(f'The Command {cmd} returned error status {return_code}')
+ if output:
+ self.log.info(output)
+
+ def create_ssh_key(self,
+ account: str,
+ expire_time: int = 300) -> str:
+ """
+ Generate a temporary SSH key and apply it to the specified account
+
+ :param account: The service account id
+ :type account: str
+ :param expire_time: The maximum amount of time before the private key is expired
+ :type expire_time: int
+ """
+ try:
+ self.log.info("Generating ssh keys...")
+ private_key_file = '/tmp/key-' + str(uuid.uuid4())
Review comment:
Paramiko is now working. Oslogin works but managing key through instant metadata still doesn't work. The key is set correctly as advised [here](https://cloud.google.com/compute/docs/instances/adding-removing-ssh-keys#instance-only) but am still yet to figure out why it's not working
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713519719
[The Workflow run](https://github.com/apache/airflow/actions/runs/319791151) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] ephraimbuddy commented on a change in pull request #9879: [WIP]Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
ephraimbuddy commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r459682916
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,327 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import pprint
+import subprocess
+import time
+import uuid
+from io import StringIO
+from typing import Any, List, Optional, Tuple, Union
+
+import paramiko
+from google.cloud.oslogin_v1 import OsLoginServiceClient
+from googleapiclient.discovery import build
+from paramiko import SSHException
+
+from airflow import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+
+class ComputeEngineSshHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version: The Compute Engine discovery API version
+ :type api_version: str
+ :param gcp_conn_id: The connection id to use when fetching connection info
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+ _oslogin_conn: Optional[OsLoginServiceClient] = None
+ _sshclient: Optional[paramiko.SSHClient] = None
+ _conn: Optional[Any] = None
+
+ def __init__(
+ self,
+ api_version: str = 'v1',
+ gcp_conn_id: str = 'google_cloud_default',
+ delegate_to: Optional[str] = None
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+
+ def get_conn(self) -> Any:
+ """
+ Retrieves connection to Google Compute Engine.
+
+ :return: Google Compute Engine services object
+ :rtype: dict
+ """
+ if not self._conn:
+ http_authorized = self._authorize()
+ self._conn = build('compute', self.api_version,
+ http=http_authorized, cache_discovery=False)
+ return self._conn
+
+ def get_oslogin_conn(self) -> OsLoginServiceClient:
+ """
+ Connects to Oslogin API
+ """
+ if not self._oslogin_conn:
+ self._oslogin_conn = OsLoginServiceClient(
+ credentials=self._get_credentials(), client_info=self.client_info
+ )
+
+ return self._oslogin_conn
+
+ def get_sshclient(self,
+ private_key_file: str,
+ hostname: str,
+ username: str
+ ) -> paramiko.SSHClient:
+ """
+ Get SSHClient
+
+ :param private_key_file: The private key filename
+ :type private_key_file: str
+ :param hostname: The hostname of the target instance
+ :type hostname: str
+ :param username: The POXIS username from the login profile of the user account
+ :type username: str
+ """
+
+ if not self._sshclient:
+ self._sshclient = paramiko.SSHClient()
+ # Default is RejectPolicy
+ # No knownhost checking since we are not storing pkey
+ self._sshclient.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+ self._sshclient.connect(hostname=hostname,
+ username=username,
+ key_filename=private_key_file,
+ look_for_keys=False)
+ return self._sshclient
+
+ def get_service_account(self) -> str:
+ """
+ Get the service account id
+ """
+ account: str = self._get_credentials().service_account_email
+ os_login_client = self.get_oslogin_conn()
+ return os_login_client.user_path(user=account)
+
+ def get_username(self, oslogin=True) -> str:
+ """
+ Get the account username
+ """
+ account = self._get_credentials().service_account_email
+ if oslogin:
+ profile = self.get_oslogin_conn().get_login_profile(name=self.get_service_account())
+ account = profile.posix_accounts[0]
+ return account.username
+ return account.split("@")[0]
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_connection_type(self, project_id: Optional[str], instance: str, zone: str,
+ hostname: Optional[str] = None,
+ use_iap_tunnel: bool = False,
+ use_internal_ip: bool = False
+ ) -> str:
+ """
+ Selects how to connect to the target instance
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param hostname: The hostname of the target instance. If provided, this would
+ be used, regardless of the setting for use_iap_tunnel and use_internal_ip
+ :type hostname: str
+ :param use_iap_tunnel: Whether to connect through IAP tunnel
+ :type use_iap_tunnel: bool
+ :param use_internal_ip: Whether to connect using internal IP
+ :type use_internal_ip: bool
+ """
+ if not hostname:
+ instance_info = self.get_conn() \
+ .instances().get(project=project_id,
+ zone=zone,
+ instance=instance).execute()
+ if use_iap_tunnel or use_internal_ip:
+ return instance_info["networkInterfaces"][0] \
+ .get("networkIP")
+ else:
+ external_ip = instance_info \
+ .get("networkInterfaces")[0] \
+ .get("accessConfigs")
+ if external_ip:
+ return external_ip[0].get("natIP")
Review comment:
The natIP is inside the accessConfigs and when external IP is removed, the accessConfigs returns an empty list
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713495028
[The Workflow run](https://github.com/apache/airflow/actions/runs/319698305) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713585989
[The Workflow run](https://github.com/apache/airflow/actions/runs/319945935) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713232463
[The Workflow run](https://github.com/apache/airflow/actions/runs/318883050) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] turbaszek commented on a change in pull request #9879: Add Compute Engine SSH hook
Posted by GitBox <gi...@apache.org>.
turbaszek commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r520028522
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,304 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import shlex
+import time
+from io import StringIO
+from typing import Any, Optional
+
+import paramiko
+from cached_property import cached_property
+from google.api_core.retry import exponential_sleep_generator
+
+from airflow import AirflowException
+from airflow.providers.google.cloud.hooks.compute import ComputeEngineHook
+from airflow.providers.google.cloud.hooks.os_login import OSLoginHook
+from airflow.providers.ssh.hooks.ssh import SSHHook
+
+
+class _GCloudAuthorizedSSHClient(paramiko.SSHClient):
+ """SSH Client that maintains the context for gcloud authorization during the connection"""
+
+ def __init__(self, google_hook, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.ssh_client = paramiko.SSHClient()
+ self.google_hook = google_hook
+ self.decorator = None
+
+ def connect(self, *args, **kwargs): # pylint: disable=signature-differs
+ self.decorator = self.google_hook.provide_authorized_gcloud()
+ self.decorator.__enter__()
+ return super().connect(*args, **kwargs)
+
+ def close(self):
+ if self.decorator:
+ self.decorator.__exit__(None, None, None)
+ self.decorator = None
+ return super().close()
+
+ def __exit__(self, type_, value, traceback):
+ if self.decorator:
+ self.decorator.__exit__(type_, value, traceback)
+ self.decorator = None
+ return super().__exit__(type_, value, traceback)
+
+
+class ComputeEngineSSHHook(SSHHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param instance_name: The name of the Compute Engine instance
+ :type instance_name: str
+ :param zone: The zone of the Compute Engine instance
+ :type zone: str
+ :param user: The name of the user on which the login attempt will be made
+ :type user: str
+ :param project_id: The project ID of the remote instance
+ :type project_id: str
+ :param gcp_conn_id: The connection id to use when fetching connection info
+ :type gcp_conn_id: str
+ :param hostname: The hostname of the target instance. If it is not passed, it will be detected
+ automatically.
+ :type hostname: str
+ :param use_iap_tunnel: Whether to connect through IAP tunnel
+ :type use_iap_tunnel: bool
+ :param use_internal_ip: Whether to connect using internal IP
+ :type use_internal_ip: bool
+ :param use_oslogin: Whether to manage keys using OsLogin API. If false,
+ keys are managed using instance metadata
+ :param expire_time: The maximum amount of time in seconds before the private key expires
+ :type expire_time: int
+ :param gcp_conn_id: The connection id to use when fetching connection information
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+
+ def __init__( # pylint: disable=too-many-arguments
+ self,
+ gcp_conn_id: str = 'google_cloud_default',
+ instance_name: Optional[str] = None,
+ zone: Optional[str] = None,
+ user: Optional[str] = 'root',
+ project_id: Optional[str] = None,
+ hostname: Optional[str] = None,
+ use_internal_ip: bool = False,
+ use_iap_tunnel: bool = False,
+ use_oslogin: bool = True,
+ expire_time: int = 300,
+ delegate_to: Optional[str] = None,
+ ) -> None:
+ # Ignore original constructor
+ # super().__init__() # pylint: disable=super-init-not-called
+ self.instance_name = instance_name
+ self.zone = zone
+ self.user = user
+ self.project_id = project_id
+ self.hostname = hostname
+ self.use_internal_ip = use_internal_ip
+ self.use_iap_tunnel = use_iap_tunnel
+ self.use_oslogin = use_oslogin
+ self.expire_time = expire_time
+ self.gcp_conn_id = gcp_conn_id
+ self.delegate_to = delegate_to
+ self._conn: Optional[Any] = None
+
+ @cached_property
+ def _oslogin_hook(self) -> OSLoginHook:
+ return OSLoginHook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to)
+
+ @cached_property
+ def _compute_hook(self) -> ComputeEngineHook:
+ return ComputeEngineHook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to)
+
+ def _load_connection_config(self):
+ def _boolify(value):
+ if isinstance(value, bool):
+ return value
+ if isinstance(value, str):
+ if value.lower() == 'false':
+ return False
+ elif value.lower() == 'true':
+ return True
+ return False
+
+ def intify(key, value, default):
+ if value is None:
+ return default
+ if isinstance(value, str) and value.strip() == '':
+ return default
+ try:
+ return int(value)
+ except ValueError:
+ raise AirflowException(
+ f"The {key} field should be a integer. "
+ f"Current value: \"{value}\" (type: {type(value)}). "
+ f"Please check the connection configuration."
+ )
+
+ conn = self.get_connection(self.gcp_conn_id)
+ if conn and conn.conn_type == "gcpssh":
+ self.instance_name = self._compute_hook._get_field( # pylint: disable=protected-access
+ "instance_name", self.instance_name
+ )
+ self.zone = self._compute_hook._get_field("zone", self.zone) # pylint: disable=protected-access
+ self.user = conn.login if conn.login else self.user
+ # self.project_id is skipped intentionally
+ self.hostname = conn.host if conn.host else self.hostname
+ self.use_internal_ip = _boolify(
+ self._compute_hook._get_field("use_internal_ip") # pylint: disable=protected-access
+ )
+ self.use_iap_tunnel = _boolify(
+ self._compute_hook._get_field("use_iap_tunnel") # pylint: disable=protected-access
+ )
+ self.use_oslogin = _boolify(
+ self._compute_hook._get_field("use_oslogin") # pylint: disable=protected-access
+ )
+ self.expire_time = intify(
+ "expire_time",
+ self._compute_hook._get_field("expire_time"), # pylint: disable=protected-access
+ self.expire_time,
+ )
+
+ def get_conn(self) -> paramiko.SSHClient:
+ """Return SSH connection."""
+ self._load_connection_config()
+ if not self.project_id:
+ self.project_id = self._compute_hook.project_id
+
+ missing_fields = [k for k in ["instance_name", "zone", "project_id"] if not getattr(self, k)]
+ if not self.instance_name or not self.zone or not self.project_id:
+ raise AirflowException(
+ f"Required parameters are missing: {missing_fields}. These parameters be passed either as "
+ "keyword parameter or as extra field in Airfow connection definition. Both are not set!"
+ )
+
+ self.log.info(
+ "Connecting to instance: instance_name=%s, user=%s, zone=%s, "
+ "use_internal_ip=%s, use_iap_tunnel=%s, use_os_login=%s",
+ self.instance_name,
+ self.user,
+ self.zone,
+ self.use_internal_ip,
+ self.use_iap_tunnel,
+ self.use_oslogin,
+ )
+ if not self.hostname:
+ hostname = self._compute_hook.get_instance_address(
+ zone=self.zone,
+ resource_id=self.instance_name,
+ project_id=self.project_id,
+ use_internal_ip=self.use_internal_ip or self.use_iap_tunnel,
+ )
+ else:
+ hostname = self.hostname
+
+ privkey, pubkey = self._generate_ssh_key(self.user)
+ if self.use_oslogin:
+ user = self._authorize_os_login(pubkey)
+ else:
+ user = self.user
+ self._authorize_compute_engine_instance_metadata(pubkey)
+
+ proxy_command = None
+ if self.use_iap_tunnel:
+ proxy_command_args = [
+ 'gcloud',
+ 'compute',
+ 'start-iap-tunnel',
+ str(self.instance_name),
+ '22',
+ '--listen-on-stdin',
+ f'--project={self.project_id}',
+ f'--zone={self.zone}',
+ '--verbosity=warning',
+ ]
+ proxy_command = " ".join(shlex.quote(arg) for arg in proxy_command_args)
+
+ sshclient = self._connect_to_instance(user, hostname, privkey, proxy_command)
+ return sshclient
+
+ def _connect_to_instance(self, user, hostname, pkey, proxy_command) -> paramiko.SSHClient:
+ self.log.info("Opening remote connection to host: username=%s, hostname=%s", user, hostname)
+ for time_to_wait in exponential_sleep_generator(initial=1, maximum=10):
+ try:
+ client = _GCloudAuthorizedSSHClient(self._compute_hook)
+ # Default is RejectPolicy
+ # No known host checking since we are not storing privatekey
+ client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+ client.connect(
+ hostname=hostname,
+ username=user,
+ pkey=pkey,
+ sock=paramiko.ProxyCommand(proxy_command) if proxy_command else None,
+ look_for_keys=False,
+ )
+ return client
+ except paramiko.SSHException:
+ if time_to_wait == 10:
Review comment:
Got it, should we create `max_wait_time` variable to keep in sync the `if` clause and generator params?
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] ephraimbuddy closed pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
ephraimbuddy closed pull request #9879:
URL: https://github.com/apache/airflow/pull/9879
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713152792
[The Workflow run](https://github.com/apache/airflow/actions/runs/318579598) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] stale[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
stale[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-691464741
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-712867002
[The Workflow run](https://github.com/apache/airflow/actions/runs/317643451) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713496038
[The Workflow run](https://github.com/apache/airflow/actions/runs/319720287) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713172612
[The Workflow run](https://github.com/apache/airflow/actions/runs/318643419) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713464503
[The Workflow run](https://github.com/apache/airflow/actions/runs/319610979) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on a change in pull request #9879: [WIP]Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r456886645
##########
File path: airflow/providers/google/cloud/hooks/compute_engine_ssh.py
##########
@@ -0,0 +1,211 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import subprocess
+import time
+import uuid
+from typing import Any, List, Optional, Union
+
+import requests
+from googleapiclient.discovery import build
+
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+SERVICE_ACCOUNT_METADATA_URL = (
+ 'http://metadata.google.internal/computeMetadata/v1/instance/'
+ 'service-accounts/default/email')
+HEADERS = {'Metadata-Flavor': 'Google'}
+
+
+class ComputeEngineSSHHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version:
+ :param gcp_conn_id:
+ :param delegate_to:
+ :param oslogin:
+ :param account:
+ """
+
+ def __init__(self,
+ api_version: str = "v1",
+ gcp_conn_id: str = "google_cloud_default",
+ delegate_to: Optional[str] = None,
+ oslogin=None,
+ account=None
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+ self.oslogin = oslogin
+ self.account = account
+
+ def get_oslogin(self) -> Any:
+ """
+ Connects to Oslogin API
+ """
+ if self.oslogin is None:
+ http_authorized = self._authorize()
+ self.oslogin = build("oslogin", self.api_version,
+ http=http_authorized, cache_discovery=False)
+ return self.oslogin
+
+ def get_service_account(self) -> str:
+ """
+ Get the service account id
+ :return:
+ """
+ if self.account is None:
+ self.account = requests.get(
+ SERVICE_ACCOUNT_METADATA_URL, headers=HEADERS).text
+ if not self.account.startswith('users/'):
+ self.account = 'users/' + self.account
+ return self.account
+
+ def get_username(self):
+ """
+ Using the OS Login API, get the POSIX username from the login profile
+ for the service account.
+
+ """
+ oslogin = self.get_oslogin()
+ account = self.get_service_account()
+ profile = oslogin.users().getLoginProfile( # pylint: disable=no-member
+ name=account).execute_in_compute()
+ username = profile.get('posixAccounts')[0].get('username')
+ return username
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_hostname(self, project_id, instance, zone, hostname=None):
+ """
+ Create the target instance hostname
+ :param project_id:
+ :param instance:
+ :param zone:
+ :param hostname:
+ :return:
+ """
+ hostname = hostname or f'{instance}.{zone}.c.{project_id}.internal'
+ return hostname
+
+ def execute_in_compute(self, cmd: Union[str, List[str]]) -> None:
+ """
+ Executes command in Compute Engine instance
+ :param cmd: The command to execute_in_compute
+ :type cmd: union[str, List[str]]
+ """
+ self.log.info('Executing command: %s', str(cmd))
+ process = subprocess.Popen(cmd)
+ output = process.communicate()[0]
+ returncode = process.returncode
+ if returncode:
+ # Error
+ self.log.error('Command returned error status %s', returncode)
+ if output:
+ self.log.info(output)
+
+ def create_ssh_key(self,
+ oslogin: Any,
+ account: str,
+ private_key_file: Optional[str] = None,
+ expire_time: int = 300) -> str:
+ """
+ Generate a temporary SSH key and apply it to the specified account
+
+ :param oslogin:
+ :param account:
+ :param private_key_file:
+ :param expire_time:
+ """
+ private_key_file = private_key_file or '/tmp/key-' + str(uuid.uuid4())
+ self.execute_in_compute(['ssh-keygen', '-t', 'rsa', '-N', '', '-f', private_key_file])
+
+ with open(private_key_file + '.pub', 'r') as original:
+ public_key = original.read().strip()
+
+ # Expiration time is in microseconds.
+ expiration = int((time.time() + expire_time) * 1000000)
+
+ body = {
+ 'key': public_key,
+ 'expirationTimeUsec': expiration,
+ }
+ oslogin.users().importSshPublicKey(parent=account, body=body).execute_in_compute()
+ return private_key_file
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def run(self,
+ cmd: str, project_id: str, instance: str, zone: str,
+ private_key_file: Optional[str] = None, expire_time: int = 300,
+ hostname: Optional[str] = None, tunnel: bool = False,
+ clean_up_private_key: bool = True):
+ """
+ Run a command on a remote system.
+
+ :param cmd:
+ :param project_id:
+ :param instance:
+ :param zone:
+ :param private_key_file:
+ :param expire_time:
+ :param hostname:
+ :param tunnel:
+ :param clean_up_private_key:
+ :return:
+ """
+
+ hostname = self.get_hostname(project_id=project_id,
+ instance=instance,
+ zone=zone,
+ hostname=zone)
+
+ account = self.get_service_account()
+ oslogin = self.get_oslogin()
+
+ private_key_file = self.create_ssh_key(oslogin=oslogin,
+ account=account,
+ private_key_file=private_key_file,
+ expire_time=expire_time)
+
+ username = self.get_username()
+
+ if tunnel and "--tunnel-through-iap" not in cmd:
+ cmd = cmd + " " + "--tunnel-through-iap"
+
+ ssh_command = [
+ 'ssh', '-i', private_key_file, '-o', 'StrictHostKeyChecking=no',
+ f'{username}@{hostname}',
+ cmd,
+ ]
+ ssh: Any = subprocess.Popen(
+ ssh_command, shell=False, stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE)
+ result = ssh.stdout.readlines()
+ if clean_up_private_key:
+ self.clean_up(private_key_file)
+ return result if result else ssh.stderr.readlines()
+
+ def clean_up(self, private_key_file: str) -> None:
+ """
+ Removes created private key file
+
+ :param private_key_file:
+ :type: str
+ """
+ self.execute_in_compute(['shred', private_key_file])
+ self.execute_in_compute(['rm', private_key_file])
Review comment:
If possible, we should avoid saving the private key to disk.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] ephraimbuddy commented on a change in pull request #9879: [WIP]Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
ephraimbuddy commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r459367980
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,327 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import pprint
+import subprocess
+import time
+import uuid
+from io import StringIO
+from typing import Any, List, Optional, Tuple, Union
+
+import paramiko
+from google.cloud.oslogin_v1 import OsLoginServiceClient
+from googleapiclient.discovery import build
+from paramiko import SSHException
+
+from airflow import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+
+class ComputeEngineSshHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version: The Compute Engine discovery API version
+ :type api_version: str
+ :param gcp_conn_id: The connection id to use when fetching connection info
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+ _oslogin_conn: Optional[OsLoginServiceClient] = None
+ _sshclient: Optional[paramiko.SSHClient] = None
+ _conn: Optional[Any] = None
+
+ def __init__(
+ self,
+ api_version: str = 'v1',
+ gcp_conn_id: str = 'google_cloud_default',
+ delegate_to: Optional[str] = None
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+
+ def get_conn(self) -> Any:
+ """
+ Retrieves connection to Google Compute Engine.
+
+ :return: Google Compute Engine services object
+ :rtype: dict
+ """
+ if not self._conn:
+ http_authorized = self._authorize()
+ self._conn = build('compute', self.api_version,
+ http=http_authorized, cache_discovery=False)
+ return self._conn
+
+ def get_oslogin_conn(self) -> OsLoginServiceClient:
+ """
+ Connects to Oslogin API
+ """
+ if not self._oslogin_conn:
+ self._oslogin_conn = OsLoginServiceClient(
+ credentials=self._get_credentials(), client_info=self.client_info
+ )
+
+ return self._oslogin_conn
+
+ def get_sshclient(self,
+ private_key_file: str,
+ hostname: str,
+ username: str
+ ) -> paramiko.SSHClient:
+ """
+ Get SSHClient
+
+ :param private_key_file: The private key filename
+ :type private_key_file: str
+ :param hostname: The hostname of the target instance
+ :type hostname: str
+ :param username: The POXIS username from the login profile of the user account
+ :type username: str
+ """
+
+ if not self._sshclient:
+ self._sshclient = paramiko.SSHClient()
+ # Default is RejectPolicy
+ # No knownhost checking since we are not storing pkey
+ self._sshclient.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+ self._sshclient.connect(hostname=hostname,
+ username=username,
+ key_filename=private_key_file,
+ look_for_keys=False)
+ return self._sshclient
+
+ def get_service_account(self) -> str:
+ """
+ Get the service account id
+ """
+ account: str = self._get_credentials().service_account_email
+ os_login_client = self.get_oslogin_conn()
+ return os_login_client.user_path(user=account)
+
+ def get_username(self, oslogin=True) -> str:
+ """
+ Get the account username
+ """
+ account = self._get_credentials().service_account_email
+ if oslogin:
+ profile = self.get_oslogin_conn().get_login_profile(name=self.get_service_account())
+ account = profile.posix_accounts[0]
+ return account.username
+ return account.split("@")[0]
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_connection_type(self, project_id: Optional[str], instance: str, zone: str,
+ hostname: Optional[str] = None,
+ use_iap_tunnel: bool = False,
+ use_internal_ip: bool = False
+ ) -> str:
+ """
+ Selects how to connect to the target instance
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param hostname: The hostname of the target instance. If provided, this would
+ be used, regardless of the setting for use_iap_tunnel and use_internal_ip
+ :type hostname: str
+ :param use_iap_tunnel: Whether to connect through IAP tunnel
+ :type use_iap_tunnel: bool
+ :param use_internal_ip: Whether to connect using internal IP
+ :type use_internal_ip: bool
+ """
+ if not hostname:
+ instance_info = self.get_conn() \
+ .instances().get(project=project_id,
+ zone=zone,
+ instance=instance).execute()
+ if use_iap_tunnel or use_internal_ip:
+ return instance_info["networkInterfaces"][0] \
+ .get("networkIP")
+ else:
+ external_ip = instance_info \
+ .get("networkInterfaces")[0] \
+ .get("accessConfigs")
+ if external_ip:
+ return external_ip[0].get("natIP")
+ else:
+ raise AirflowException("The target instance does not have external IP,"
+ "consider specifying use_internal_ip or use_iap_tunnel")
+
+ return hostname
+
+ def execute_local(self, cmd: Union[str, List[str]]) -> None:
+ """
+ Executes command in local machine
+
+ :param cmd: The command to execute
+ :type cmd: Union[str, List[str]]
+ """
+ self.log.info('Executing command: %s', str(cmd))
+ process = subprocess.Popen(cmd)
+ output = process.communicate()[0]
+ return_code = process.returncode
+ if return_code:
+ # Error
+ raise AirflowException(f'The Command {cmd} returned error status {return_code}')
+ if output:
+ self.log.info(output)
+
+ def _create_ssh_key(self,
+ instance: str,
+ zone: str,
+ project_id: Optional[str],
+ expire_time: int = 300,
+ use_oslogin: bool = True) -> str:
+ """
+ Generate a temporary SSH key and apply it to the specified account
+
+ :param expire_time: The maximum amount of time before the private key is expired
+ :type expire_time: int
+ :param use_oslogin: Whether to use OsLogin to manage key
+ """
+ username = self.get_username(oslogin=use_oslogin)
+ try:
+ self.log.info("Generating ssh keys...")
+ # private_key_file = StringIO()
+ # private_key_obj = paramiko.RSAKey.generate(2048)
+ # private_key_obj.write_private_key(private_key_file)
+
+ # public_key = f"{private_key_obj.get_name()} {private_key_obj.get_base64()}" \
+ # f" {username}"
+ private_key_file = '/tmp/key-' + str(uuid.uuid4())
+ self.execute_local(['ssh-keygen', '-t', 'rsa', '-N', '', '-f', private_key_file])
+
+ with open(private_key_file + '.pub', 'r') as original:
+ public_key = original.read().strip()
+
+ except (IOError, SSHException) as err:
+ raise AirflowException(f"Error encountered creating ssh keys, {err}")
+
+ if use_oslogin:
+ # Expiration time is in microseconds.
+ expiration = int((time.time() + expire_time) * 1000000)
+
+ body = {
+ "key": public_key,
+ "expiration_time_usec": expiration
+ }
+ os_login_client = self.get_oslogin_conn()
+ os_login_client.import_ssh_public_key(
+ parent=self.get_service_account(),
+ ssh_public_key=body
+ )
+ else:
+ # using instance metadata to manage keys
+
+ instance_info = self.get_conn().instances().get(
+ project=project_id,
+ instance=instance,
+ zone=zone).execute()
+
+ new_keys = [username + ":" + public_key + "\n"]
+ metadata = instance_info['metadata']
+ items = metadata.get("items")
+ if items:
+ key_vals = items[0].get("value")
+ for val in key_vals:
+ new_keys.append(val + "\n")
+ new_dict = {"key": "ssh-keys", "value": str(new_keys)}
+ metadata.update(new_dict)
+ self.get_conn().instances().setMetadata(
+ project=project_id,
+ zone=zone,
+ instance=instance,
+ body=metadata
+ ).execute()
Review comment:
This is not being set
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on a change in pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r519460158
##########
File path: airflow/providers/ssh/operators/ssh.py
##########
@@ -89,7 +90,8 @@ def execute(self, context) -> Union[bytes, str, bool]:
self.log.info(
"ssh_hook is not provided or invalid. " "Trying ssh_conn_id to create SSHHook."
)
- self.ssh_hook = SSHHook(ssh_conn_id=self.ssh_conn_id, timeout=self.timeout)
+ self.ssh_hook = BaseHook.get_connection(self.ssh_conn_id)
Review comment:
I have withdrawn this change. I will deal with it in a separate PR.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on a change in pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r460595488
##########
File path: airflow/providers/google/cloud/operators/compute_ssh.py
##########
@@ -0,0 +1,105 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from typing import Any, Dict, Optional
+
+from airflow.models import BaseOperator
+from airflow.providers.google.cloud.hooks.compute_ssh import ComputeEngineSshHook
+from airflow.utils.decorators import apply_defaults
+
+
+class ComputeEngineSSHOperator(BaseOperator):
+ """
+ Executes commands on a remote instance
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:ComputeEngineSSHOperator`
+
+ :param instance: The name of the remote instance
+ :type instance: str
+ :param zone: The zone of the remote instance
+ :type zone: str
+ :param cmd: The command to execute on the remote instance
+ :type cmd: str
+ :param project_id: The project ID of the remote instance
+ :type project_id: str
+ :param gcp_conn_id: The connection id to use when fetching connection info
+ :type gcp_conn_id: str
+ :param api_version: The Compute Engine discovery API version
+ :type api_version: str
+ :param hostname: The hostname of the target instance
+ :type hostname: str
+ :param use_iap_tunnel: Whether to connect through IAP tunnel
+ :type use_iap_tunnel: bool
+ :param use_internal_ip: Whether to connect using internal IP
+ :type use_internal_ip: bool
+ :param use_oslogin: Whether to manage keys using OsLogin API. If false,
+ keys are managed using instance metadata
+ :param expire_time: The maximum amount of time in seconds before the private key expires
+ :type expire_time: int
+ """
+ # [START remote_instance_start_template_fields]
+ template_fields = ('project_id', 'zone', 'instance', 'gcp_conn_id', 'api_version')
+
+ # [END remote_instance_start_template_fields]
+
+ @apply_defaults
+ def __init__(self, # pylint: disable=too-many-arguments
+ instance: str,
+ zone: str,
+ cmd: str,
+ project_id: Optional[str] = None,
+ gcp_conn_id: str = 'google_cloud_default',
+ api_version: str = "v1",
+ hostname: Optional[str] = None,
+ use_internal_ip: bool = False,
+ use_iap_tunnel: bool = False,
+ use_oslogin: bool = True,
+ expiry_time: int = 300,
+ *args: Any,
Review comment:
```suggestion
```
See: https://github.com/apache/airflow/issues/9942
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713459650
[The Workflow run](https://github.com/apache/airflow/actions/runs/319597659) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713451033
[The Workflow run](https://github.com/apache/airflow/actions/runs/319568661) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713165349
[The Workflow run](https://github.com/apache/airflow/actions/runs/318606200) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on a change in pull request #9879: [WIP]Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r458453184
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,250 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import subprocess
+import time
+import uuid
+from typing import Any, List, Optional, Union
+
+import paramiko
+from googleapiclient.discovery import build
+from paramiko import SSHException
+
+from airflow import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+
+class ComputeEngineSshHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version: The Oslogin discovery API version
+ :type api_version: str
+ :param gcp_conn_id: The connection id to use when fetching connection info
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+
+ def __init__(self,
+ api_version: str = "v1",
+ gcp_conn_id: str = "google_cloud_default",
+ delegate_to: Optional[str] = None
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+ self.conn = None
+
+ def get_conn(self) -> Any:
+ """
+ Connects to Oslogin API
+ """
+ if self.conn is None:
+ http_authorized = self._authorize()
+ self.conn = build("oslogin", self.api_version,
+ http=http_authorized, cache_discovery=False)
+ return self.conn
+
+ def get_sshclient(self,
+ private_key_file: str,
+ hostname: str,
+ username: str,
+ strict_host_checking: bool = True
+ ) -> paramiko.SSHClient:
+ """
+ Get SSHClient
+
+ :param private_key_file: The filename of the private key
+ :type private_key_file: str
+ :param hostname: The hostname of the target instance
+ :type hostname: str
+ :param username: The POXIS username from the login profile of the user account
+ :type username: str
+ :param strict_host_checking: Whether to verify key
+ :type strict_host_checking: bool
+ """
+
+ self.log.info("Creating remote connection to host")
+ client = paramiko.SSHClient()
+ if not strict_host_checking:
+ self.log.warning('No Host Key Verification. This wont protect '
+ 'against Man-In-The-Middle attacks')
+ # Default is RejectPolicy
+ client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+ client.connect(hostname=hostname,
+ username=username,
+ key_filename=private_key_file)
+ return client
+
+ def get_service_account(self) -> str:
+ """
+ Get the service account id
+ """
+ account: str = self._get_credentials().service_account_email
+ if not account.startswith('users/'):
+ account = 'users/' + account
+ return account
+
+ def get_username(self) -> str:
+ """
+ Using the OS Login API, get the POSIX username from the login profile
+ for the service account.
+
+ """
+ conn = self.get_conn()
+ account = self.get_service_account()
+ profile = conn.users().getLoginProfile(name=account).execute_in_compute() # pylint: disable=no-member
+ username: str = profile.get('posixAccounts')[0].get('username')
+ return username
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_hostname(self, project_id: Optional[str], instance: str, zone: str,
+ hostname: Optional[str] = None) -> str:
+ """
+ Create the target instance hostname
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param hostname: The hostname of the target instance. If not provided,
+ it is generated from zone, instance name and project id
+ :type hostname: str
+ """
+ if not hostname:
+ hostname = f'{instance}.{zone}.c.{project_id}.internal'
+ return hostname
+
+ def execute_local(self, cmd: Union[str, List[str]]) -> None:
+ """
+ Executes command in local machine
+
+ :param cmd: The command to execute
+ :type cmd: Union[str, List[str]]
+ """
+ self.log.info('Executing command: %s', str(cmd))
+ process = subprocess.Popen(cmd)
+ output = process.communicate()[0]
+ return_code = process.returncode
+ if return_code:
+ # Error
+ raise AirflowException(f'The Command {cmd} returned error status {return_code}')
+ if output:
+ self.log.info(output)
+
+ def create_ssh_key(self,
+ account: str,
+ expire_time: int = 300) -> str:
+ """
+ Generate a temporary SSH key and apply it to the specified account
+
+ :param account: The service account id
+ :type account: str
+ :param expire_time: The maximum amount of time before the private key is expired
+ :type expire_time: int
+ """
+ try:
+ self.log.info("Generating ssh keys...")
+ private_key_file = '/tmp/key-' + str(uuid.uuid4())
+ private_key_obj = paramiko.RSAKey.generate(2048)
+ private_key_obj.write_private_key_file(private_key_file)
+ except (IOError, SSHException) as err:
+ raise AirflowException(f"Error encountered creating ssh keys, {err}")
+
+ with open(private_key_file + '.pub', 'r') as original:
+ public_key = original.read().strip()
+
+ # Expiration time is in microseconds.
+ expiration = int((time.time() + expire_time) * 1000000)
+
+ body = {
+ 'key': public_key,
+ 'expirationTimeUsec': expiration,
+ }
+ self.get_conn(). \
+ users().importSshPublicKey( # pylint: disable=no-member
+ parent=account, body=body).execute_in_compute()
+ return private_key_file
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def run(self,
+ cmd: str, instance: str, zone: str,
+ project_id: Optional[str] = None,
+ expire_time: int = 300,
+ hostname: Optional[str] = None,
+ strict_host_key_checking: bool = True
+ ) -> Any:
+ """
+ Execute a command in a remote instance
+
+ :param cmd: The command to execute on the instance
+ :type cmd: str
+ :param instance: The name of the instance the command will be executed on
+ :type instance: str
+ :param zone: The zone where the instance is located
+ :type zone: str
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param expire_time: The maximum amount of time before the private key is expired
+ :type expire_time: int
+ :param hostname: The hostname of the instance.
+ :type hostname: str
+ :param strict_host_key_checking: Whether to verify the host key
+ :type strict_host_key_checking: bool
+ """
+ hostname = self.get_hostname(project_id=project_id,
+ instance=instance,
+ zone=zone,
+ hostname=hostname)
+
+ account = self.get_service_account()
+
+ username = self.get_username()
+
+ private_key_file = self.create_ssh_key(account=account,
Review comment:
It would be great if we also supported authorizations with instance metadata.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on a change in pull request #9879: [WIP]Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r459457981
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,327 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import pprint
+import subprocess
+import time
+import uuid
+from io import StringIO
+from typing import Any, List, Optional, Tuple, Union
+
+import paramiko
+from google.cloud.oslogin_v1 import OsLoginServiceClient
+from googleapiclient.discovery import build
+from paramiko import SSHException
+
+from airflow import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+
+class ComputeEngineSshHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version: The Compute Engine discovery API version
+ :type api_version: str
+ :param gcp_conn_id: The connection id to use when fetching connection info
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+ _oslogin_conn: Optional[OsLoginServiceClient] = None
Review comment:
This creates a class attribute, not an instance attribute
https://mypy.readthedocs.io/en/stable/class_basics.html
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713530350
[The Workflow run](https://github.com/apache/airflow/actions/runs/319829745) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713429932
[The Workflow run](https://github.com/apache/airflow/actions/runs/319505196) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713583555
[The Workflow run](https://github.com/apache/airflow/actions/runs/319919029) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713023877
[The Workflow run](https://github.com/apache/airflow/actions/runs/318146469) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj edited a comment on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj edited a comment on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-664808821
@ephraimbuddy Can you try this terraform module? This should automatically configure the necessary resources.
https://github.com/terraform-google-modules/terraform-google-bastion-host/tree/master/examples/iap_tunneling
https://github.com/terraform-google-modules/terraform-google-bastion-host
If you still run into problems, let me know and I will configure the VM for you.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713727160
[The Workflow run](https://github.com/apache/airflow/actions/runs/320242963) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713775801
[The Workflow run](https://github.com/apache/airflow/actions/runs/320410508) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713471940
We have a ticket already: https://support.github.com/ticket/personal/0/877699
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-712996672
[The Workflow run](https://github.com/apache/airflow/actions/runs/318001343) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on a change in pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r460596997
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,378 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import subprocess
+import time
+from io import StringIO
+from typing import Any, List, Optional, Union
+
+import paramiko
+from google.cloud.oslogin_v1 import OsLoginServiceClient
+from googleapiclient.discovery import build
+from paramiko import AuthenticationException, BadHostKeyException, SSHException
+
+from airflow import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+
+class ComputeEngineSshHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version: The Compute Engine discovery API version
+ :type api_version: str
+ :param gcp_conn_id: The connection id to use when fetching connection information
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+ _conn: Optional[Any] = None
+
+ def __init__(
+ self,
+ api_version: str = 'v1',
+ gcp_conn_id: str = 'google_cloud_default',
+ delegate_to: Optional[str] = None
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+ self._oslogin_conn: Optional[OsLoginServiceClient] = None
+
+ def get_conn(self) -> Any:
+ """
+ Retrieves connection to Google Compute Engine.
+
+ :return: Google Compute Engine services object
+ :rtype: dict
+ """
+ if not self._conn:
+ http_authorized = self._authorize()
+ self._conn = build('compute', self.api_version,
+ http=http_authorized, cache_discovery=False)
+ return self._conn
+
+ def get_oslogin_conn(self) -> OsLoginServiceClient:
+ """
+ Connects to Oslogin API
+ """
+ if not self._oslogin_conn:
+ self._oslogin_conn = OsLoginServiceClient(
+ credentials=self._get_credentials(), client_info=self.client_info
+ )
+
+ return self._oslogin_conn
+
+ def get_sshclient(self,
+ pkey: str,
+ hostname: str,
+ username: str
+ ) -> paramiko.SSHClient:
+ """
+ Get SSHClient
+
+ :param pkey: The private key
+ :type pkey: str
+ :param hostname: The hostname of the target instance
+ :type hostname: str
+ :param username: The username from the login profile of the user account
+ :type username: str
+ """
+
+ client = paramiko.SSHClient()
+ # Default is RejectPolicy
+ # No knownhost checking since we are not storing pkey
+ client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+ client.connect(hostname=hostname,
+ username=username,
+ pkey=pkey,
+ look_for_keys=False)
+ return client
+
+ def get_service_account(self) -> str:
+ """
+ Get the service account id
+ """
+ account: str = self._get_credentials().service_account_email
+ os_login_client = self.get_oslogin_conn()
+ return os_login_client.user_path(user=account)
+
+ def get_username(self, oslogin=True) -> str:
+ """
+ Get the account username
+
+ :param oslogin: Whether to use OsLogin
+ :type oslogin: bool
+ """
+ account = self._get_credentials().service_account_email
+ if oslogin:
+ profile = self.get_oslogin_conn().get_login_profile(name=self.get_service_account())
+ account = profile.posix_accounts[0]
+ return account.username
+ return account.split("@")[0]
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_connection_type(self, project_id: Optional[str], instance: str, zone: str,
+ hostname: Optional[str] = None,
+ use_iap_tunnel: bool = False,
+ use_internal_ip: bool = False
+ ) -> str:
+ """
+ Selects how to connect to the target instance
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param hostname: The hostname of the target instance. If provided, this would
+ be used, regardless of the setting for use_iap_tunnel and use_internal_ip
+ :type hostname: str
+ :param use_iap_tunnel: Whether to connect through IAP tunnel
+ :type use_iap_tunnel: bool
+ :param use_internal_ip: Whether to connect using internal IP
+ :type use_internal_ip: bool
+ """
+ if not hostname:
+ instance_info = self.instance_info(project_id=project_id,
+ zone=zone,
+ instance=instance)
+ if use_iap_tunnel or use_internal_ip:
+ return instance_info["networkInterfaces"][0] \
+ .get("networkIP")
+ else:
+ access_config = instance_info \
+ .get("networkInterfaces")[0] \
+ .get("accessConfigs")
+ if access_config:
+ return access_config[0].get("natIP")
+ else:
+ raise AirflowException("The target instance does not have external IP,"
+ "consider specifying use_internal_ip or use_iap_tunnel")
+ return hostname
+
+ def execute(self, cmd: Union[str, List[str]]) -> None:
+ """
+ Executes command in local machine
+
+ :param cmd: The command to execute
+ :type cmd: Union[str, List[str]]
+ """
+ self.log.info('Executing command: %s', str(cmd))
+ process = subprocess.Popen(cmd)
+ output = process.communicate()[0]
+ returncode = process.returncode
+ if returncode:
+ # Error
+ self.log.error('Command returned error status %s', returncode)
+ if output:
+ self.log.info(output)
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def instance_info(self, project_id, zone, instance):
+ """
+ Used to get instance information
+
+ :param project_id: The project ID where the instance is located
+ :type project_id: str
+ :param zone: The zone of the instance
+ :type zone: str
+ :param instance: The name of the instance
+ :type instance: str
+ """
+ instance_info = self.get_conn().instances().get( # pylint: disable=no-member
+ project=project_id,
+ instance=instance,
+ zone=zone).execute()
+ return instance_info
+
+ def _update_metadata(self, keys, metadata):
+ """
+ Updates ssh keys on metadata dictionary
+
+ :param keys: The public key to set on the metadata
+ :type keys: str
+ :param metadata: A dictionary containing instance metadata
+ :type metadata: dict
+ """
+ items = metadata.get("items")
+ for item in items:
+ if item.get("key") == "ssh-keys":
+ keys += item["value"]
+ item['value'] = keys
+ break
+ else:
+ new_dict = dict(key='ssh-keys', value=keys)
+ metadata['items'] = [new_dict]
+ return metadata
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def set_sshkey_metadata(self, project_id, zone, instance, username, pubkey):
+ """
+ Sets ssh key on instance metadata
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param username: The logged in account username
+ :type username: str
+ :param pubkey: The public key to set on the instance
+ :type pubkey: str
+ """
+ instance_info = self.instance_info(project_id=project_id,
+ zone=zone,
+ instance=instance)
+ keys = username + ":" + pubkey + "\n"
+ metadata = instance_info['metadata']
+ metadata = self._update_metadata(keys, metadata)
+ self.get_conn().instances().setMetadata( # pylint: disable=no-member
+ project=project_id,
+ zone=zone,
+ instance=instance,
+ body=metadata
+ ).execute()
+ time.sleep(10) # TODO improve on this. The status should be Done before proceeding
+
+ def set_oslogin_key(self, pubkey, expire_time):
+ """
+ Sets public key on instance os login
+
+ :param pubkey: The public key to set on the OsLogin
+ :type pubkey: str
+ :param expire_time: The amount of time before the key expires in seconds
+ :type expire_time: int
+ """
+ expiration = int((time.time() + expire_time) * 1000000)
+ body = {
+ "key": pubkey,
+ "expiration_time_usec": expiration
+ }
+ os_login_client = self.get_oslogin_conn()
+ os_login_client.import_ssh_public_key(
+ parent=self.get_service_account(),
+ ssh_public_key=body
+ )
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def _create_ssh_key(self,
+ instance: str,
+ zone: str,
+ project_id: Optional[str] = None,
+ expire_time: int = 300,
+ use_oslogin: bool = True) -> Any:
+ """
+ Generate a temporary SSH key and apply it to the specified account
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param expire_time: The maximum amount of time before the private key is expired
+ :type expire_time: int
+ :param use_oslogin: Whether to use OsLogin to manage key
+ :type use_oslogin: int
+ """
+ username = self.get_username(oslogin=use_oslogin)
+ try:
+ self.log.info("Generating ssh keys...")
+ pkey_file = StringIO()
+ pkey_obj = paramiko.RSAKey.generate(2048)
+ pkey_obj.write_private_key(pkey_file)
+ pubkey = f"{pkey_obj.get_name()} {pkey_obj.get_base64()} {username}"
+ except (IOError, SSHException) as err:
+ raise AirflowException(f"Error encountered creating ssh keys, {err}")
+
+ if use_oslogin:
+ # Expiration time is in microseconds.
+ self.set_oslogin_key(pubkey, expire_time)
+ else:
+ # using instance metadata to manage keys
+ self.set_sshkey_metadata(project_id=project_id,
+ zone=zone,
+ instance=instance,
+ username=username,
+ pubkey=pubkey)
+ return pkey_obj
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def run(self,
+ cmd: str, instance: str, zone: str,
+ project_id: Optional[str] = None,
+ expire_time: int = 300,
+ hostname: Optional[str] = None,
+ use_internal_ip: bool = False,
+ use_iap_tunnel: bool = False,
+ use_oslogin: bool = True
+ ) -> Any:
+ """
+ Execute a command in a remote instance
+
+ :param cmd: The command to execute on the instance
+ :type cmd: str
+ :param instance: The name of the instance the command will be executed on
+ :type instance: str
+ :param zone: The zone where the instance is located
+ :type zone: str
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param expire_time: The maximum amount of time before the private key is expired
+ :type expire_time: int
+ :param hostname: The hostname of the instance.
+ :type hostname: str
+ :param use_iap_tunnel: Whether to connect through IAP tunnel
+ :type use_iap_tunnel: bool
+ :param use_internal_ip: Whether to connect using internal IP
+ :type use_internal_ip: bool
+ :param use_oslogin: Whether to manage keys using OsLogin API. If false,
+ keys are managed using instance metadata
+ """
+ hostname = self.get_connection_type(project_id=project_id,
+ instance=instance,
+ zone=zone,
+ hostname=hostname,
+ use_iap_tunnel=use_iap_tunnel,
+ use_internal_ip=use_internal_ip
+ )
+
+ username = self.get_username(oslogin=use_oslogin)
+
+ self.log.info("Creating remote connection to host")
+ pkey = self._create_ssh_key(
+ instance=instance,
+ project_id=project_id,
+ zone=zone,
+ expire_time=expire_time,
+ use_oslogin=use_oslogin
+ )
+ sshclient = self.get_sshclient(pkey=pkey,
Review comment:
There is one problem if internal IP addresses are used. Two servers can use the same IP address if they are on two different networks. We should make sure that we have connected to the expected server.
```
metadata_id_url = 'http://metadata.google.internal/computeMetadata/v1/instance/id'
remote_command = f'[ `curl "{metadata_id_url}" -H "Metadata-Flavor: Google" -q` = {instance_id} ] || exit 42'
```
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713132542
[The Workflow run](https://github.com/apache/airflow/actions/runs/318510560) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713195279
[The Workflow run](https://github.com/apache/airflow/actions/runs/318731586) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-712935623
[The Workflow run](https://github.com/apache/airflow/actions/runs/317913958) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713207893
[The Workflow run](https://github.com/apache/airflow/actions/runs/318772674) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-664808821
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713224503
[The Workflow run](https://github.com/apache/airflow/actions/runs/318845166) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713138721
[The Workflow run](https://github.com/apache/airflow/actions/runs/318531276) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] ephraimbuddy commented on a change in pull request #9879: [WIP]Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
ephraimbuddy commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r459682916
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,327 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import pprint
+import subprocess
+import time
+import uuid
+from io import StringIO
+from typing import Any, List, Optional, Tuple, Union
+
+import paramiko
+from google.cloud.oslogin_v1 import OsLoginServiceClient
+from googleapiclient.discovery import build
+from paramiko import SSHException
+
+from airflow import AirflowException
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+
+class ComputeEngineSshHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version: The Compute Engine discovery API version
+ :type api_version: str
+ :param gcp_conn_id: The connection id to use when fetching connection info
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+ _oslogin_conn: Optional[OsLoginServiceClient] = None
+ _sshclient: Optional[paramiko.SSHClient] = None
+ _conn: Optional[Any] = None
+
+ def __init__(
+ self,
+ api_version: str = 'v1',
+ gcp_conn_id: str = 'google_cloud_default',
+ delegate_to: Optional[str] = None
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+
+ def get_conn(self) -> Any:
+ """
+ Retrieves connection to Google Compute Engine.
+
+ :return: Google Compute Engine services object
+ :rtype: dict
+ """
+ if not self._conn:
+ http_authorized = self._authorize()
+ self._conn = build('compute', self.api_version,
+ http=http_authorized, cache_discovery=False)
+ return self._conn
+
+ def get_oslogin_conn(self) -> OsLoginServiceClient:
+ """
+ Connects to Oslogin API
+ """
+ if not self._oslogin_conn:
+ self._oslogin_conn = OsLoginServiceClient(
+ credentials=self._get_credentials(), client_info=self.client_info
+ )
+
+ return self._oslogin_conn
+
+ def get_sshclient(self,
+ private_key_file: str,
+ hostname: str,
+ username: str
+ ) -> paramiko.SSHClient:
+ """
+ Get SSHClient
+
+ :param private_key_file: The private key filename
+ :type private_key_file: str
+ :param hostname: The hostname of the target instance
+ :type hostname: str
+ :param username: The POXIS username from the login profile of the user account
+ :type username: str
+ """
+
+ if not self._sshclient:
+ self._sshclient = paramiko.SSHClient()
+ # Default is RejectPolicy
+ # No knownhost checking since we are not storing pkey
+ self._sshclient.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+ self._sshclient.connect(hostname=hostname,
+ username=username,
+ key_filename=private_key_file,
+ look_for_keys=False)
+ return self._sshclient
+
+ def get_service_account(self) -> str:
+ """
+ Get the service account id
+ """
+ account: str = self._get_credentials().service_account_email
+ os_login_client = self.get_oslogin_conn()
+ return os_login_client.user_path(user=account)
+
+ def get_username(self, oslogin=True) -> str:
+ """
+ Get the account username
+ """
+ account = self._get_credentials().service_account_email
+ if oslogin:
+ profile = self.get_oslogin_conn().get_login_profile(name=self.get_service_account())
+ account = profile.posix_accounts[0]
+ return account.username
+ return account.split("@")[0]
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_connection_type(self, project_id: Optional[str], instance: str, zone: str,
+ hostname: Optional[str] = None,
+ use_iap_tunnel: bool = False,
+ use_internal_ip: bool = False
+ ) -> str:
+ """
+ Selects how to connect to the target instance
+
+ :param project_id: The project ID of the instance
+ :type project_id: Optional[str]
+ :param instance: The name of the target instance to connect to
+ :type instance: str
+ :param zone: The zone of the target instance
+ :param zone: str
+ :param hostname: The hostname of the target instance. If provided, this would
+ be used, regardless of the setting for use_iap_tunnel and use_internal_ip
+ :type hostname: str
+ :param use_iap_tunnel: Whether to connect through IAP tunnel
+ :type use_iap_tunnel: bool
+ :param use_internal_ip: Whether to connect using internal IP
+ :type use_internal_ip: bool
+ """
+ if not hostname:
+ instance_info = self.get_conn() \
+ .instances().get(project=project_id,
+ zone=zone,
+ instance=instance).execute()
+ if use_iap_tunnel or use_internal_ip:
+ return instance_info["networkInterfaces"][0] \
+ .get("networkIP")
+ else:
+ external_ip = instance_info \
+ .get("networkInterfaces")[0] \
+ .get("accessConfigs")
+ if external_ip:
+ return external_ip[0].get("natIP")
Review comment:
The natIP is inside the accessConfigs and when external IP is removed, the accessConfigs returns an empty list. I have renamed it correctly
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on a change in pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r519459692
##########
File path: airflow/providers/google/cloud/hooks/compute_ssh.py
##########
@@ -0,0 +1,304 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+import shlex
+import time
+from io import StringIO
+from typing import Any, Optional
+
+import paramiko
+from cached_property import cached_property
+from google.api_core.retry import exponential_sleep_generator
+
+from airflow import AirflowException
+from airflow.providers.google.cloud.hooks.compute import ComputeEngineHook
+from airflow.providers.google.cloud.hooks.os_login import OSLoginHook
+from airflow.providers.ssh.hooks.ssh import SSHHook
+
+
+class _GCloudAuthorizedSSHClient(paramiko.SSHClient):
+ """SSH Client that maintains the context for gcloud authorization during the connection"""
+
+ def __init__(self, google_hook, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self.ssh_client = paramiko.SSHClient()
+ self.google_hook = google_hook
+ self.decorator = None
+
+ def connect(self, *args, **kwargs): # pylint: disable=signature-differs
+ self.decorator = self.google_hook.provide_authorized_gcloud()
+ self.decorator.__enter__()
+ return super().connect(*args, **kwargs)
+
+ def close(self):
+ if self.decorator:
+ self.decorator.__exit__(None, None, None)
+ self.decorator = None
+ return super().close()
+
+ def __exit__(self, type_, value, traceback):
+ if self.decorator:
+ self.decorator.__exit__(type_, value, traceback)
+ self.decorator = None
+ return super().__exit__(type_, value, traceback)
+
+
+class ComputeEngineSSHHook(SSHHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param instance_name: The name of the Compute Engine instance
+ :type instance_name: str
+ :param zone: The zone of the Compute Engine instance
+ :type zone: str
+ :param user: The name of the user on which the login attempt will be made
+ :type user: str
+ :param project_id: The project ID of the remote instance
+ :type project_id: str
+ :param gcp_conn_id: The connection id to use when fetching connection info
+ :type gcp_conn_id: str
+ :param hostname: The hostname of the target instance. If it is not passed, it will be detected
+ automatically.
+ :type hostname: str
+ :param use_iap_tunnel: Whether to connect through IAP tunnel
+ :type use_iap_tunnel: bool
+ :param use_internal_ip: Whether to connect using internal IP
+ :type use_internal_ip: bool
+ :param use_oslogin: Whether to manage keys using OsLogin API. If false,
+ keys are managed using instance metadata
+ :param expire_time: The maximum amount of time in seconds before the private key expires
+ :type expire_time: int
+ :param gcp_conn_id: The connection id to use when fetching connection information
+ :type gcp_conn_id: str
+ :param delegate_to: The account to impersonate, if any.
+ For this to work, the service account making the request must have
+ domain-wide delegation enabled.
+ :type delegate_to: str
+ """
+
+ def __init__( # pylint: disable=too-many-arguments
+ self,
+ gcp_conn_id: str = 'google_cloud_default',
+ instance_name: Optional[str] = None,
+ zone: Optional[str] = None,
+ user: Optional[str] = 'root',
+ project_id: Optional[str] = None,
+ hostname: Optional[str] = None,
+ use_internal_ip: bool = False,
+ use_iap_tunnel: bool = False,
+ use_oslogin: bool = True,
+ expire_time: int = 300,
+ delegate_to: Optional[str] = None,
+ ) -> None:
+ # Ignore original constructor
+ # super().__init__() # pylint: disable=super-init-not-called
+ self.instance_name = instance_name
+ self.zone = zone
+ self.user = user
+ self.project_id = project_id
+ self.hostname = hostname
+ self.use_internal_ip = use_internal_ip
+ self.use_iap_tunnel = use_iap_tunnel
+ self.use_oslogin = use_oslogin
+ self.expire_time = expire_time
+ self.gcp_conn_id = gcp_conn_id
+ self.delegate_to = delegate_to
+ self._conn: Optional[Any] = None
+
+ @cached_property
+ def _oslogin_hook(self) -> OSLoginHook:
+ return OSLoginHook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to)
+
+ @cached_property
+ def _compute_hook(self) -> ComputeEngineHook:
+ return ComputeEngineHook(gcp_conn_id=self.gcp_conn_id, delegate_to=self.delegate_to)
+
+ def _load_connection_config(self):
+ def _boolify(value):
+ if isinstance(value, bool):
+ return value
+ if isinstance(value, str):
+ if value.lower() == 'false':
+ return False
+ elif value.lower() == 'true':
+ return True
+ return False
+
+ def intify(key, value, default):
+ if value is None:
+ return default
+ if isinstance(value, str) and value.strip() == '':
+ return default
+ try:
+ return int(value)
+ except ValueError:
+ raise AirflowException(
+ f"The {key} field should be a integer. "
+ f"Current value: \"{value}\" (type: {type(value)}). "
+ f"Please check the connection configuration."
+ )
+
+ conn = self.get_connection(self.gcp_conn_id)
+ if conn and conn.conn_type == "gcpssh":
+ self.instance_name = self._compute_hook._get_field( # pylint: disable=protected-access
+ "instance_name", self.instance_name
+ )
+ self.zone = self._compute_hook._get_field("zone", self.zone) # pylint: disable=protected-access
+ self.user = conn.login if conn.login else self.user
+ # self.project_id is skipped intentionally
+ self.hostname = conn.host if conn.host else self.hostname
+ self.use_internal_ip = _boolify(
+ self._compute_hook._get_field("use_internal_ip") # pylint: disable=protected-access
+ )
+ self.use_iap_tunnel = _boolify(
+ self._compute_hook._get_field("use_iap_tunnel") # pylint: disable=protected-access
+ )
+ self.use_oslogin = _boolify(
+ self._compute_hook._get_field("use_oslogin") # pylint: disable=protected-access
+ )
+ self.expire_time = intify(
+ "expire_time",
+ self._compute_hook._get_field("expire_time"), # pylint: disable=protected-access
+ self.expire_time,
+ )
+
+ def get_conn(self) -> paramiko.SSHClient:
+ """Return SSH connection."""
+ self._load_connection_config()
+ if not self.project_id:
+ self.project_id = self._compute_hook.project_id
+
+ missing_fields = [k for k in ["instance_name", "zone", "project_id"] if not getattr(self, k)]
+ if not self.instance_name or not self.zone or not self.project_id:
+ raise AirflowException(
+ f"Required parameters are missing: {missing_fields}. These parameters be passed either as "
+ "keyword parameter or as extra field in Airfow connection definition. Both are not set!"
+ )
+
+ self.log.info(
+ "Connecting to instance: instance_name=%s, user=%s, zone=%s, "
+ "use_internal_ip=%s, use_iap_tunnel=%s, use_os_login=%s",
+ self.instance_name,
+ self.user,
+ self.zone,
+ self.use_internal_ip,
+ self.use_iap_tunnel,
+ self.use_oslogin,
+ )
+ if not self.hostname:
+ hostname = self._compute_hook.get_instance_address(
+ zone=self.zone,
+ resource_id=self.instance_name,
+ project_id=self.project_id,
+ use_internal_ip=self.use_internal_ip or self.use_iap_tunnel,
+ )
+ else:
+ hostname = self.hostname
+
+ privkey, pubkey = self._generate_ssh_key(self.user)
+ if self.use_oslogin:
+ user = self._authorize_os_login(pubkey)
+ else:
+ user = self.user
+ self._authorize_compute_engine_instance_metadata(pubkey)
+
+ proxy_command = None
+ if self.use_iap_tunnel:
+ proxy_command_args = [
+ 'gcloud',
+ 'compute',
+ 'start-iap-tunnel',
+ str(self.instance_name),
+ '22',
+ '--listen-on-stdin',
+ f'--project={self.project_id}',
+ f'--zone={self.zone}',
+ '--verbosity=warning',
+ ]
+ proxy_command = " ".join(shlex.quote(arg) for arg in proxy_command_args)
+
+ sshclient = self._connect_to_instance(user, hostname, privkey, proxy_command)
+ return sshclient
+
+ def _connect_to_instance(self, user, hostname, pkey, proxy_command) -> paramiko.SSHClient:
+ self.log.info("Opening remote connection to host: username=%s, hostname=%s", user, hostname)
+ for time_to_wait in exponential_sleep_generator(initial=1, maximum=10):
+ try:
+ client = _GCloudAuthorizedSSHClient(self._compute_hook)
+ # Default is RejectPolicy
+ # No known host checking since we are not storing privatekey
+ client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+ client.connect(
+ hostname=hostname,
+ username=user,
+ pkey=pkey,
+ sock=paramiko.ProxyCommand(proxy_command) if proxy_command else None,
+ look_for_keys=False,
+ )
+ return client
+ except paramiko.SSHException:
+ if time_to_wait == 10:
Review comment:
exponential_sleep_generator is an infinite generator, so we need to manually check the end condition.
```
>>> pprint(list(zip(exponential_sleep_generator(1, 10), range(1, 30))))
[(0.3548877463516451, 1),
(3.2179802309727807, 2),
(6.267463665198731, 3),
(10, 4),
(10, 5),
(10, 6),
(10, 7),
(10, 8),
(10, 9),
(10, 10),
(10, 11),
(10, 12),
(10, 13),
(10, 14),
(10, 15),
(10, 16),
(10, 17),
(10, 18),
(10, 19),
(10, 20),
(10, 21),
(10, 22),
(10, 23),
(10, 24),
(10, 25),
(10, 26),
(10, 27),
(10, 28),
(10, 29)]
```
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Add Compute Engine SSH hook
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-724198042
The PR needs to run all tests because it modifies core of Airflow! Please rebase it to latest master or ask committer to re-run 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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713257139
[The Workflow run](https://github.com/apache/airflow/actions/runs/318978177) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] mik-laj commented on a change in pull request #9879: [WIP]Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
mik-laj commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r456886544
##########
File path: airflow/providers/google/cloud/hooks/compute_engine_ssh.py
##########
@@ -0,0 +1,211 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import subprocess
+import time
+import uuid
+from typing import Any, List, Optional, Union
+
+import requests
+from googleapiclient.discovery import build
+
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+SERVICE_ACCOUNT_METADATA_URL = (
+ 'http://metadata.google.internal/computeMetadata/v1/instance/'
+ 'service-accounts/default/email')
+HEADERS = {'Metadata-Flavor': 'Google'}
+
+
+class ComputeEngineSSHHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version:
+ :param gcp_conn_id:
+ :param delegate_to:
+ :param oslogin:
+ :param account:
+ """
+
+ def __init__(self,
+ api_version: str = "v1",
+ gcp_conn_id: str = "google_cloud_default",
+ delegate_to: Optional[str] = None,
+ oslogin=None,
+ account=None
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+ self.oslogin = oslogin
+ self.account = account
+
+ def get_oslogin(self) -> Any:
+ """
+ Connects to Oslogin API
+ """
+ if self.oslogin is None:
+ http_authorized = self._authorize()
+ self.oslogin = build("oslogin", self.api_version,
+ http=http_authorized, cache_discovery=False)
+ return self.oslogin
+
+ def get_service_account(self) -> str:
+ """
+ Get the service account id
+ :return:
+ """
+ if self.account is None:
+ self.account = requests.get(
+ SERVICE_ACCOUNT_METADATA_URL, headers=HEADERS).text
+ if not self.account.startswith('users/'):
+ self.account = 'users/' + self.account
+ return self.account
+
+ def get_username(self):
+ """
+ Using the OS Login API, get the POSIX username from the login profile
+ for the service account.
+
+ """
+ oslogin = self.get_oslogin()
+ account = self.get_service_account()
+ profile = oslogin.users().getLoginProfile( # pylint: disable=no-member
+ name=account).execute_in_compute()
+ username = profile.get('posixAccounts')[0].get('username')
+ return username
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_hostname(self, project_id, instance, zone, hostname=None):
+ """
+ Create the target instance hostname
+ :param project_id:
+ :param instance:
+ :param zone:
+ :param hostname:
+ :return:
+ """
+ hostname = hostname or f'{instance}.{zone}.c.{project_id}.internal'
+ return hostname
+
+ def execute_in_compute(self, cmd: Union[str, List[str]]) -> None:
+ """
+ Executes command in Compute Engine instance
+ :param cmd: The command to execute_in_compute
+ :type cmd: union[str, List[str]]
+ """
+ self.log.info('Executing command: %s', str(cmd))
+ process = subprocess.Popen(cmd)
+ output = process.communicate()[0]
+ returncode = process.returncode
+ if returncode:
+ # Error
+ self.log.error('Command returned error status %s', returncode)
+ if output:
+ self.log.info(output)
+
+ def create_ssh_key(self,
+ oslogin: Any,
+ account: str,
+ private_key_file: Optional[str] = None,
+ expire_time: int = 300) -> str:
+ """
+ Generate a temporary SSH key and apply it to the specified account
+
+ :param oslogin:
+ :param account:
+ :param private_key_file:
+ :param expire_time:
+ """
+ private_key_file = private_key_file or '/tmp/key-' + str(uuid.uuid4())
+ self.execute_in_compute(['ssh-keygen', '-t', 'rsa', '-N', '', '-f', private_key_file])
+
+ with open(private_key_file + '.pub', 'r') as original:
+ public_key = original.read().strip()
+
+ # Expiration time is in microseconds.
+ expiration = int((time.time() + expire_time) * 1000000)
+
+ body = {
+ 'key': public_key,
+ 'expirationTimeUsec': expiration,
+ }
+ oslogin.users().importSshPublicKey(parent=account, body=body).execute_in_compute()
+ return private_key_file
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def run(self,
+ cmd: str, project_id: str, instance: str, zone: str,
+ private_key_file: Optional[str] = None, expire_time: int = 300,
+ hostname: Optional[str] = None, tunnel: bool = False,
+ clean_up_private_key: bool = True):
+ """
+ Run a command on a remote system.
+
+ :param cmd:
+ :param project_id:
+ :param instance:
+ :param zone:
+ :param private_key_file:
+ :param expire_time:
+ :param hostname:
+ :param tunnel:
+ :param clean_up_private_key:
+ :return:
+ """
+
+ hostname = self.get_hostname(project_id=project_id,
+ instance=instance,
+ zone=zone,
+ hostname=zone)
+
+ account = self.get_service_account()
+ oslogin = self.get_oslogin()
+
+ private_key_file = self.create_ssh_key(oslogin=oslogin,
+ account=account,
+ private_key_file=private_key_file,
+ expire_time=expire_time)
+
+ username = self.get_username()
+
+ if tunnel and "--tunnel-through-iap" not in cmd:
+ cmd = cmd + " " + "--tunnel-through-iap"
+
+ ssh_command = [
Review comment:
We should try to do it with paramiko. We can then use more advanced features. Although I know it will be a bit more difficult then.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] ephraimbuddy commented on a change in pull request #9879: [WIP]Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
ephraimbuddy commented on a change in pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#discussion_r456888648
##########
File path: airflow/providers/google/cloud/hooks/compute_engine_ssh.py
##########
@@ -0,0 +1,211 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import subprocess
+import time
+import uuid
+from typing import Any, List, Optional, Union
+
+import requests
+from googleapiclient.discovery import build
+
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+SERVICE_ACCOUNT_METADATA_URL = (
+ 'http://metadata.google.internal/computeMetadata/v1/instance/'
+ 'service-accounts/default/email')
+HEADERS = {'Metadata-Flavor': 'Google'}
+
+
+class ComputeEngineSSHHook(GoogleBaseHook):
+ """
+ Hook to connect to a remote instance in compute engine
+
+ :param api_version:
+ :param gcp_conn_id:
+ :param delegate_to:
+ :param oslogin:
+ :param account:
+ """
+
+ def __init__(self,
+ api_version: str = "v1",
+ gcp_conn_id: str = "google_cloud_default",
+ delegate_to: Optional[str] = None,
+ oslogin=None,
+ account=None
+ ) -> None:
+ super().__init__(gcp_conn_id, delegate_to)
+ self.api_version = api_version
+ self.oslogin = oslogin
+ self.account = account
+
+ def get_oslogin(self) -> Any:
+ """
+ Connects to Oslogin API
+ """
+ if self.oslogin is None:
+ http_authorized = self._authorize()
+ self.oslogin = build("oslogin", self.api_version,
+ http=http_authorized, cache_discovery=False)
+ return self.oslogin
+
+ def get_service_account(self) -> str:
+ """
+ Get the service account id
+ :return:
+ """
+ if self.account is None:
+ self.account = requests.get(
+ SERVICE_ACCOUNT_METADATA_URL, headers=HEADERS).text
+ if not self.account.startswith('users/'):
+ self.account = 'users/' + self.account
+ return self.account
+
+ def get_username(self):
+ """
+ Using the OS Login API, get the POSIX username from the login profile
+ for the service account.
+
+ """
+ oslogin = self.get_oslogin()
+ account = self.get_service_account()
+ profile = oslogin.users().getLoginProfile( # pylint: disable=no-member
+ name=account).execute_in_compute()
+ username = profile.get('posixAccounts')[0].get('username')
+ return username
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def get_hostname(self, project_id, instance, zone, hostname=None):
+ """
+ Create the target instance hostname
+ :param project_id:
+ :param instance:
+ :param zone:
+ :param hostname:
+ :return:
+ """
+ hostname = hostname or f'{instance}.{zone}.c.{project_id}.internal'
+ return hostname
+
+ def execute_in_compute(self, cmd: Union[str, List[str]]) -> None:
+ """
+ Executes command in Compute Engine instance
+ :param cmd: The command to execute_in_compute
+ :type cmd: union[str, List[str]]
+ """
+ self.log.info('Executing command: %s', str(cmd))
+ process = subprocess.Popen(cmd)
+ output = process.communicate()[0]
+ returncode = process.returncode
+ if returncode:
+ # Error
+ self.log.error('Command returned error status %s', returncode)
+ if output:
+ self.log.info(output)
+
+ def create_ssh_key(self,
+ oslogin: Any,
+ account: str,
+ private_key_file: Optional[str] = None,
+ expire_time: int = 300) -> str:
+ """
+ Generate a temporary SSH key and apply it to the specified account
+
+ :param oslogin:
+ :param account:
+ :param private_key_file:
+ :param expire_time:
+ """
+ private_key_file = private_key_file or '/tmp/key-' + str(uuid.uuid4())
+ self.execute_in_compute(['ssh-keygen', '-t', 'rsa', '-N', '', '-f', private_key_file])
+
+ with open(private_key_file + '.pub', 'r') as original:
+ public_key = original.read().strip()
+
+ # Expiration time is in microseconds.
+ expiration = int((time.time() + expire_time) * 1000000)
+
+ body = {
+ 'key': public_key,
+ 'expirationTimeUsec': expiration,
+ }
+ oslogin.users().importSshPublicKey(parent=account, body=body).execute_in_compute()
+ return private_key_file
+
+ @GoogleBaseHook.fallback_to_default_project_id
+ def run(self,
+ cmd: str, project_id: str, instance: str, zone: str,
+ private_key_file: Optional[str] = None, expire_time: int = 300,
+ hostname: Optional[str] = None, tunnel: bool = False,
+ clean_up_private_key: bool = True):
+ """
+ Run a command on a remote system.
+
+ :param cmd:
+ :param project_id:
+ :param instance:
+ :param zone:
+ :param private_key_file:
+ :param expire_time:
+ :param hostname:
+ :param tunnel:
+ :param clean_up_private_key:
+ :return:
+ """
+
+ hostname = self.get_hostname(project_id=project_id,
+ instance=instance,
+ zone=zone,
+ hostname=zone)
+
+ account = self.get_service_account()
+ oslogin = self.get_oslogin()
+
+ private_key_file = self.create_ssh_key(oslogin=oslogin,
+ account=account,
+ private_key_file=private_key_file,
+ expire_time=expire_time)
+
+ username = self.get_username()
+
+ if tunnel and "--tunnel-through-iap" not in cmd:
+ cmd = cmd + " " + "--tunnel-through-iap"
+
+ ssh_command = [
+ 'ssh', '-i', private_key_file, '-o', 'StrictHostKeyChecking=no',
+ f'{username}@{hostname}',
+ cmd,
+ ]
+ ssh: Any = subprocess.Popen(
+ ssh_command, shell=False, stdout=subprocess.PIPE,
+ stderr=subprocess.PIPE)
+ result = ssh.stdout.readlines()
+ if clean_up_private_key:
+ self.clean_up(private_key_file)
+ return result if result else ssh.stderr.readlines()
+
+ def clean_up(self, private_key_file: str) -> None:
+ """
+ Removes created private key file
+
+ :param private_key_file:
+ :type: str
+ """
+ self.execute_in_compute(['shred', private_key_file])
+ self.execute_in_compute(['rm', private_key_file])
Review comment:
It is actually deleted even when the user didn't instruct for it to be removed. Once the expire time has elapsed, the key is deleted even when the clean_up_private_key is False but I will check if it can work without saving to disk.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-713146180
[The Workflow run](https://github.com/apache/airflow/actions/runs/318556155) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org
[GitHub] [airflow] github-actions[bot] commented on pull request #9879: Implement Compute Engine SSH hook and operator
Posted by GitBox <gi...@apache.org>.
github-actions[bot] commented on pull request #9879:
URL: https://github.com/apache/airflow/pull/9879#issuecomment-712897959
[The Workflow run](https://github.com/apache/airflow/actions/runs/317793254) is cancelling this PR. It has some failed jobs matching ^Pylint$,^Static checks$,^Build docs$,^Spell check docs$,^Backport packages$,^Checks: Helm tests$,^Test OpenAPI*.
----------------------------------------------------------------
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:
users@infra.apache.org