mage-ai 0.9.44__py3-none-any.whl → 0.9.46__py3-none-any.whl

mage_ai/services/ssh/aws/emr/models.py

@@ -0,0 +1,326 @@
+import json
+import os
+import socket
+from typing import Dict
+
+from sshtunnel import SSHTunnelForwarder
+
+from mage_ai.data_preparation.models.project import Project
+from mage_ai.services.compute.aws.models import Cluster
+from mage_ai.services.compute.constants import SSH_PORT
+from mage_ai.services.compute.models import ComputeService
+from mage_ai.services.spark.api.constants import SPARK_UI_HOST
+from mage_ai.services.ssh.aws.emr.constants import SSH_DEFAULTS
+from mage_ai.services.ssh.aws.emr.utils import file_path, should_tunnel
+from mage_ai.shared.hash import merge_dict
+from mage_ai.shared.utils import is_port_in_use
+
+
+class SSHTunnel:
+    _instance = None
+    _tunnel = None
+    ec2_key_path = None
+    master_public_dns_name = None
+    spark_ui_host_local = None
+    spark_ui_host_remote = None
+    spark_ui_port_local = None
+    spark_ui_port_remote = None
+    ssh_username = None
+
+    def __new__(
+        cls,
+        ec2_key_path: str = None,
+        master_public_dns_name: str = None,
+        spark_ui_host_local: str = None,
+        spark_ui_host_remote: str = None,
+        spark_ui_port_local: int = None,
+        spark_ui_port_remote: int = None,
+        ssh_username: str = None,
+    ):
+        if cls._instance is None:
+            if ec2_key_path and master_public_dns_name:
+                cls._instance = super(SSHTunnel, cls).__new__(cls)
+
+                """
+                ssh \
+                    -i ~/.ssh/mage-emr-2023.pem \
+                    -L [local host]:[local port]:[remote host]:[remote port] \
+                    hadoop@ec2-34-217-92-37.us-west-2.compute.amazonaws.com
+                """
+                cls._instance.ec2_key_path = ec2_key_path
+                cls._instance.master_public_dns_name = master_public_dns_name
+                cls._instance.spark_ui_host_local = spark_ui_host_local
+                cls._instance.spark_ui_host_remote = spark_ui_host_remote
+                cls._instance.spark_ui_port_local = spark_ui_port_local
+                cls._instance.spark_ui_port_remote = spark_ui_port_remote
+                cls._instance.ssh_username = ssh_username or SSH_DEFAULTS['ssh_username']
+
+                cls._instance._tunnel = SSHTunnelForwarder(
+                    cls._instance.master_public_dns_name,
+                    local_bind_address=(
+                        cls._instance.spark_ui_host_local or SSH_DEFAULTS['spark_ui_host_local'],
+                        get_port(
+                            cls._instance.spark_ui_port_local or
+                            SSH_DEFAULTS['spark_ui_port_local']
+                        ),
+                    ),
+                    remote_bind_address=(
+                        cls._instance.spark_ui_host_remote or SSH_DEFAULTS['spark_ui_host_remote'],
+                        cls._instance.spark_ui_port_remote or SSH_DEFAULTS['spark_ui_port_remote'],
+                    ),
+                    ssh_pkey=cls._instance.ec2_key_path,
+                    ssh_username=cls._instance.ssh_username,
+                )
+
+        return cls._instance
+
+    @property
+    def tunnel(self) -> SSHTunnelForwarder:
+        if self._instance:
+            return self._instance._tunnel
+
+    def connect(self) -> bool:
+        if not self.precheck_access():
+            return
+
+        if not self.is_active():
+            self.__start()
+
+        return self.is_active()
+
+    def reconnect(self) -> bool:
+        if not self.precheck_access():
+            return False
+
+        if self.is_active():
+            self.__restart()
+        else:
+            self.__start()
+
+        return self.is_active()
+
+    def stop(self) -> None:
+        if not self.precheck_access():
+            return
+
+        if self.tunnel:
+            try:
+                self.tunnel.stop()
+            except Exception as err:
+                print(f'[WARNING] AWS EMR SSHTunnel: {err}')
+
+    def close(self) -> None:
+        if not self.precheck_access():
+            return
+
+        if self.tunnel:
+            self.tunnel.stop()
+            self._instance = None
+
+            absolute_file_path = file_path()
+            if os.path.exists(absolute_file_path):
+                os.remove(absolute_file_path)
+
+    def connection_details(self) -> Dict:
+        is_active = self.is_active()
+
+        return dict(
+            address=self.tunnel.local_bind_address if is_active else None,
+            host=SPARK_UI_HOST if is_active else None,
+            port=self.tunnel.local_bind_port if is_active else None,
+        )
+
+    def to_dict(self) -> Dict:
+        return merge_dict(self.connection_details(), dict(
+            active=self.is_active(),
+            ec2_key_path=self.ec2_key_path,
+            master_public_dns_name=self.master_public_dns_name,
+            spark_ui_host_local=self.spark_ui_host_local,
+            spark_ui_host_remote=self.spark_ui_host_remote,
+            spark_ui_port_local=self.spark_ui_port_local,
+            spark_ui_port_remote=self.spark_ui_port_remote,
+            ssh_username=self.ssh_username,
+        ))
+
+    def is_active(self, raise_error: bool = False) -> bool:
+        if not self.precheck_access():
+            return False
+
+        if not self.tunnel:
+            return False
+
+        if raise_error:
+
+            self.tunnel._check_is_started()
+
+        return self.tunnel.is_active
+
+    def precheck_access(self, raise_error: bool = False) -> bool:
+        if not self._instance:
+            return False
+
+        try:
+            test_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            test_socket.settimeout(5)
+            test_socket.connect((self._instance.master_public_dns_name, SSH_PORT))
+
+            return True
+        except Exception as err:
+            if raise_error:
+                raise err
+            print(f'[WARNING] AWS EMR SSHTunnel precheck access: {err}')
+
+        return False
+
+    def __restart(self) -> None:
+        if self.tunnel:
+            try:
+                self.tunnel.restart()
+            except Exception as err:
+                print(f'[WARNING] AWS EMR SSHTunnel restart: {err}')
+
+    def __start(self) -> None:
+        if self.tunnel:
+            try:
+                self.tunnel.start()
+            except Exception as err:
+                print(f'[WARNING] AWS EMR SSHTunnel start: {err}')
+
+
+def create_tunnel(
+    clean_up_on_failure: bool = False,
+    cluster: Cluster = None,
+    compute_service: ComputeService = None,
+    connect: bool = False,
+    fresh: bool = False,
+    project: Project = None,
+    reconnect: bool = False,
+    spark_ui_host_local: str = None,
+    spark_ui_host_remote: str = None,
+    spark_ui_port_local: int = None,
+    spark_ui_port_remote: int = None,
+    ssh_username: str = None,
+    stop: bool = False,
+) -> SSHTunnel:
+    cluster_from_cache = False
+
+    absolute_file_path = file_path()
+    if os.path.exists(absolute_file_path):
+        if fresh:
+            os.remove(absolute_file_path)
+        else:
+            with open(absolute_file_path, 'r') as f:
+                text = f.read()
+                if text:
+                    commands = json.loads(text) or {}
+
+                    cluster = cluster or commands.get('cluster')
+                    if cluster and isinstance(cluster, dict):
+                        cluster = Cluster.load(**cluster)
+                        cluster_from_cache = True
+
+                    connect = connect or commands.get('connect')
+                    reconnect = reconnect or commands.get('reconnect')
+                    spark_ui_host_local = spark_ui_host_local or commands.get(
+                        'spark_ui_host_local',
+                    )
+                    spark_ui_host_remote = spark_ui_host_remote or commands.get(
+                        'spark_ui_host_remote',
+                    )
+                    spark_ui_port_local = spark_ui_port_local or commands.get(
+                        'spark_ui_port_local',
+                    )
+                    spark_ui_port_remote = spark_ui_port_remote or commands.get(
+                        'spark_ui_port_remote',
+                    )
+                    ssh_username = ssh_username or commands.get('ssh_username')
+                    stop = stop or commands.get('stop')
+
+    if not project:
+        project = Project()
+
+    if not should_tunnel(
+        ignore_active_kernel=True,
+        project=project,
+    ):
+        return close_tunnel()
+
+    if not compute_service:
+        compute_service = ComputeService.build(project)
+
+    if not compute_service:
+        return close_tunnel()
+
+    if not cluster:
+        cluster = compute_service.active_cluster()
+
+    if not cluster:
+        return close_tunnel()
+
+    if cluster_from_cache:
+        try:
+            cluster_server = compute_service.get_cluster_details(cluster.id)
+        except Exception as err:
+            print(f'[WARNING] AWS EMR create tunnel: {err}')
+            return
+
+        if not cluster_server or cluster_server.invalid:
+            return
+        active = cluster.active
+        cluster = Cluster.load(**cluster_server.to_dict())
+        cluster.active = active
+
+    try:
+        tunnel = SSHTunnel(
+            project.emr_config.get('ec2_key_path'),
+            cluster.master_public_dns_name,
+            spark_ui_host_local=spark_ui_host_local,
+            spark_ui_host_remote=spark_ui_host_remote,
+            spark_ui_port_local=spark_ui_port_local,
+            spark_ui_port_remote=spark_ui_port_remote,
+            ssh_username=ssh_username,
+        )
+
+        if connect:
+            tunnel.connect()
+
+        if reconnect:
+            tunnel.reconnect()
+
+        if stop:
+            tunnel.stop()
+    except Exception as err:
+        if clean_up_on_failure:
+            if os.path.exists(absolute_file_path):
+                os.remove(absolute_file_path)
+        raise err
+
+    # This method is invoked when the server restarts.
+    # When the server restarts, there is no active cluster set.
+    # We need to set it again.
+    if compute_service and cluster:
+        compute_service.update_cluster(cluster.id, dict(active=cluster.active))
+
+    return tunnel
+
+
+def close_tunnel() -> None:
+    tunnel = SSHTunnel()
+    if tunnel:
+        tunnel.close()
+
+
+def get_port(port: int) -> int:
+    local_port = port
+    max_local_port = int(local_port) + 100
+
+    while is_port_in_use(local_port):
+        if local_port > max_local_port:
+            raise Exception(
+                f'Cannot find an open port from {port} to '
+                f'{max_local_port} for SSH tunnel. '
+                'Please clear running processes to free up available ports.',
+            )
+        local_port += 1
+
+    return local_port

mage_ai/services/ssh/aws/emr/utils.py

@@ -0,0 +1,151 @@
+import json
+import os
+from typing import Dict
+
+import simplejson
+
+from mage_ai.data_preparation.models.project import Project
+from mage_ai.data_preparation.models.project.constants import FeatureUUID
+from mage_ai.server.kernels import KernelName
+from mage_ai.services.compute.models import ComputeService
+from mage_ai.services.spark.constants import ComputeServiceUUID
+from mage_ai.services.spark.utils import get_compute_service
+from mage_ai.settings.repo import get_repo_path
+from mage_ai.shared.parsers import encode_complex
+
+
+def file_path() -> str:
+    return os.path.join(
+        get_repo_path(),
+        '.ssh_tunnel',
+        'aws_emr.json',
+    )
+
+
+def should_tunnel(
+    ignore_active_kernel: bool = False,
+    kernel_name: KernelName = None,
+    project: Project = None,
+    raise_error: bool = False,
+) -> bool:
+    error_message = None
+    if not project:
+        project = Project()
+
+    if not project.is_feature_enabled(FeatureUUID.COMPUTE_MANAGEMENT):
+        error_message = 'Compute management feature isn’t enabled. '\
+            'Turn the feature on in the project’s settings.'
+
+        if raise_error:
+            raise Exception(error_message)
+        else:
+            return False
+
+    if ComputeServiceUUID.AWS_EMR != get_compute_service(
+        project.repo_config,
+        ignore_active_kernel=ignore_active_kernel,
+        kernel_name=kernel_name,
+    ):
+        error_message = 'SSH tunnel is only supported for the AWS EMR compute service.'
+
+        if raise_error:
+            raise Exception(error_message)
+        else:
+            return False
+
+    if not project.emr_config:
+        error_message = 'Project has incomplete configuration for AWS EMR.'
+
+        if raise_error:
+            raise Exception(error_message)
+        else:
+            return False
+
+    ec2_key_name = project.emr_config.get('ec2_key_name')
+    ec2_key_path = project.emr_config.get('ec2_key_path')
+
+    if not ec2_key_name or not ec2_key_path:
+        error_message = 'SSH tunnel requires an EC2 key name and EC2 key path ' \
+            'in the project’s settings.'
+
+        if raise_error:
+            raise Exception(error_message)
+        else:
+            return False
+
+    if not os.path.exists(ec2_key_path):
+        error_message = f'EC2 key pair {ec2_key_path} doesn’t exist.'
+
+        if raise_error:
+            raise Exception(error_message)
+
+        return False
+
+    return True
+
+
+def tunnel(
+    ignore_active_kernel: bool = False,
+    kernel_name: KernelName = None,
+    reconnect: bool = False,
+    spark_ui_host_local: str = None,
+    spark_ui_host_remote: str = None,
+    spark_ui_port_local: str = None,
+    spark_ui_port_remote: str = None,
+    ssh_username: str = None,
+    stop: bool = False,
+    validate_conditions: bool = False,
+) -> None:
+    project = Project()
+
+    if validate_conditions and not should_tunnel(
+        ignore_active_kernel=ignore_active_kernel,
+        kernel_name=kernel_name,
+        project=project,
+    ):
+        return
+
+    data = dict(
+        connect=not reconnect,
+        reconnect=reconnect,
+        spark_ui_host_local=spark_ui_host_local,
+        spark_ui_host_remote=spark_ui_host_remote,
+        spark_ui_port_local=spark_ui_port_local,
+        spark_ui_port_remote=spark_ui_port_remote,
+        ssh_username=ssh_username,
+        stop=stop,
+    )
+
+    if project and \
+            project.emr_config and \
+            project.is_feature_enabled(FeatureUUID.COMPUTE_MANAGEMENT):
+
+        compute_service = ComputeService.build(project)
+        if compute_service and ComputeServiceUUID.AWS_EMR == compute_service.uuid:
+            cluster = compute_service.active_cluster()
+            if cluster:
+                data['cluster'] = cluster.to_dict()
+
+    absolute_file_path = file_path()
+    os.makedirs(os.path.dirname(absolute_file_path), exist_ok=True)
+
+    with open(absolute_file_path, 'w') as f:
+        f.write(simplejson.dumps(
+            data,
+            default=encode_complex,
+            ignore_nan=True,
+        ))
+
+
+def cluster_info_from_tunnel() -> Dict:
+    cluster = None
+
+    absolute_file_path = file_path()
+    if os.path.exists(absolute_file_path):
+        with open(absolute_file_path, 'r') as f:
+            text = f.read()
+            if text:
+                commands = json.loads(text) or {}
+                cluster = commands.get('cluster')
+
+    return cluster

mage_ai/settings/__init__.py

@@ -1,5 +1,7 @@
 import os
 
+from .secret_generation import generate_jwt_secret
+
 # If you add a new environment variable, make sure to check if it should be added to
 # the `MAGE_SETTINGS_ENVIRONMENT_VARIABLES` list at the bottom of this file. Also, update
 # the environment variable documentation at docs/development/variables/environment-variables.mdx
@@ -94,7 +96,10 @@ REQUESTS_BASE_PATH = os.getenv('MAGE_REQUESTS_BASE_PATH', BASE_PATH)
 # Routes base path is used to configure the base path for the backend routes. Defaults
 # to the MAGE_BASE_PATH environment variable.
 ROUTES_BASE_PATH = os.getenv('MAGE_ROUTES_BASE_PATH', BASE_PATH)
-
+# Used for OAUTH
+JWT_SECRET = os.getenv('JWT_SECRET', 'materia')
+# Used for generating download tokens
+JWT_DOWNLOAD_SECRET = os.getenv('JWT_DOWNLOAD_SECRET', generate_jwt_secret())
 # Sets the trigger interval of the scheduler to a numeric value, in seconds
 # Determines how often the scheduler gets invoked
 try:
@@ -125,7 +130,17 @@ MAGE_SETTINGS_ENVIRONMENT_VARIABLES = [
     'SENTRY_TRACES_SAMPLE_RATE',
     'MAGE_PUBLIC_HOST',
     'ACTIVE_DIRECTORY_DIRECTORY_ID',
+    'ACTIVE_DIRECTORY_CLIENT_ID',
+    'ACTIVE_DIRECTORY_CLIENT_SECRET',
     'SCHEDULER_TRIGGER_INTERVAL',
     'REQUIRE_USER_PERMISSIONS',
     'ENABLE_PROMETHEUS',
+    'OKTA_DOMAIN_URL',
+    'OKTA_CLIENT_ID',
+    'OKTA_CLIENT_SECRET',
+    'GOOGLE_CLIENT_ID',
+    'GOOGLE_CLIENT_SECRET',
+    'GHE_CLIENT_ID',
+    'GHE_CLIENT_SECRET',
+    'GHE_HOSTNAME',
 ]

mage_ai/settings/secret_generation.py

@@ -0,0 +1,7 @@
+import secrets
+import string
+
+
+def generate_jwt_secret(length=64):
+    characters = string.ascii_letters + string.digits
+    return ''.join(secrets.choice(characters) for _ in range(length))

mage_ai/settings/sso.py

@@ -0,0 +1,20 @@
+import os
+
+# Environment variable settings used to set up authentication for single
+# sign on.
+
+# Okta settings
+
+OKTA_DOMAIN_URL = os.getenv('OKTA_DOMAIN_URL')
+OKTA_CLIENT_ID = os.getenv('OKTA_CLIENT_ID')
+OKTA_CLIENT_SECRET = os.getenv('OKTA_CLIENT_SECRET')
+
+# Google settings
+
+GOOGLE_CLIENT_ID = os.getenv('GOOGLE_CLIENT_ID')
+GOOGLE_CLIENT_SECRET = os.getenv('GOOGLE_CLIENT_SECRET')
+
+# Active directory settings
+
+ACTIVE_DIRECTORY_CLIENT_ID = os.getenv('ACTIVE_DIRECTORY_CLIENT_ID')
+ACTIVE_DIRECTORY_CLIENT_SECRET = os.getenv('ACTIVE_DIRECTORY_CLIENT_SECRET')

mage_ai/shared/hash.py

@@ -106,7 +106,7 @@ def index_by(func, arr):
     return obj
 
 
-def merge_dict(a, b):
+def merge_dict(a: Dict, b: Dict) -> Dict:
     if a:
         c = a.copy()
     else:
@@ -130,3 +130,19 @@ def replace_dict_nan_value(d):
 
 def get_safe_value(data: Dict, key: str, default_value):
     return data.get(key, default_value) if data else default_value
+
+
+def set_value(obj: Dict, keys: List[str], value) -> Dict:
+    if len(keys) >= 2:
+        for idx in range(len(keys)):
+            keys_init = keys[:idx]
+            if len(keys_init) >= 1:
+                set_value(obj, keys_init, dig(obj, keys_init) or {})
+
+    results = dict(__obj_to_set_value=obj, __value=value)
+
+    key = ''.join(f"['{key}']" for key in keys)
+    expression = f'__obj_to_set_value{key} = __value'
+    exec(expression, results)
+
+    return results['__obj_to_set_value']