sagemaker-studio-dataengineering-sessions 1.0.2__tar.gz

sagemaker_studio_dataengineering_sessions-1.0.2/.gitignore

@@ -0,0 +1,18 @@
+# Peru artifact directories
+/build
+/private
+/.hatch
+/external-distribution
+
+# Python
+*.egg-info/
+__pycache__/
+*__pycache__/
+.coverage
+.ruff_cache
+.pytest_cache
+.mypy_cache
+.venv
+/htmlcov
+
+.idea/

sagemaker_studio_dataengineering_sessions-1.0.2/PKG-INFO

@@ -0,0 +1,114 @@
+Metadata-Version: 2.4
+Name: sagemaker-studio-dataengineering-sessions
+Version: 1.0.2
+Requires-Python: >=3.11
+Requires-Dist: amazon-sagemaker-sql-execution<1,>=0.0.7
+Requires-Dist: boto3
+Requires-Dist: botocore
+Requires-Dist: ipython
+Requires-Dist: pandas
+Requires-Dist: panel
+Requires-Dist: pytest
+Requires-Dist: pytest-cov
+Requires-Dist: sparkmagic
+Requires-Dist: sqlparse>=0.4.0
+Description-Content-Type: text/markdown
+
+# SageMakerStudioDataEngineeringSessions
+
+SageMaker Unified Studio Data Engineering Sessions
+
+This pacakge depends on SageMaker Unified Studio environment, if you are using SageMaker Unified Studio, see [AWS Doc](https://docs.aws.amazon.com/sagemaker-unified-studio/latest/userguide/what-is-sagemaker-unified-studio.html) for guidance.
+
+This package contains functionality to support SageMaker Unified Studio connecting to various AWS Compute including EMR/EMR Serverless/Glue/Redshift etc. 
+
+It is utilizing [ipython magics](https://ipython.readthedocs.io/en/stable/interactive/magics.html) and [AWS DataZone Connections](https://docs.aws.amazon.com/datazone/latest/APIReference/API_ListConnections.html) to achieve the following features.
+
+## Features
+
+- Connect to remote compute
+- Execute Spark code in remote compute in Python/Scala
+- Execute SQL queries in remote compute
+- Send local variables to remote compute
+
+
+## How to setup
+
+If you are using SageMaker Unifed Studio, you can skip this part, SageMaker Unifed Studio already set up the package.
+
+This package contains various Jupyter Magics to achieve its functionality.
+
+To load these magics, make sure you have iPython config file generated. If not, you could run `ipython profile create`, then a file with path `~/.ipython/profile_default/ipython_config.py` should be generated
+
+Then you will need to add the following line in the end of that config file
+
+```
+c.InteractiveShellApp.extensions.extend(['sagemaker_studio_dataengineering_sessions.sagemaker_connection_magic'])
+```
+
+Once that is finished, you could restart the ipython kernel and run `%help` to see a list of supported magics
+
+## Examples
+
+
+To connect to remote compute, a DataZone Connection is required, you could create it via [CreateConnection API](https://docs.aws.amazon.com/datazone/latest/APIReference/API_CreateConnection.html), Let's say there's an existing connection called project.spark. 
+
+### Supported Connection Type:
+
+- IAM
+- SPARK
+- REDSHIFT
+- ATHENA
+
+### Connect to remote compute and Execute Spark Code in Python
+The following example will connect to AWS Glue Interactive session and run the spark code in Glue.
+
+```
+%%pyspark project.spark
+
+import sys
+import boto3
+from awsglue.utils import getResolvedOptions
+from pyspark.context import SparkContext
+from pyspark.sql import SparkSession
+from pyspark.sql.functions import col
+
+args = getResolvedOptions(sys.argv, ["redshift_url", "redshift_iam_role", "redshift_tempdir","redshift_jdbc_iam_url"])
+print(f"{args}")
+
+sc = SparkContext.getOrCreate()
+spark = SparkSession(sc)
+
+df = spark.read.csv(f"s3://sagemaker-example-files-prod-{boto3.session.Session().region_name}/datasets/tabular/dirty-titanic/", header=True)
+df.show(5, truncate=False)
+df.printSchema()
+
+df.createOrReplaceTempView("df_sql_tempview")
+```
+
+### Execute Spark Code in Scala
+The following example will connect to AWS Glue Interactive session and run the spark code in Scala.
+
+```
+%%scalaspark project.spark
+val dfScala = spark.sql("SELECT count(0) FROM df_sql_tempview")
+dfScala.show()
+```
+
+### Execute SQL query in remote compute
+The following example will connect to AWS Glue Interactive session and run the spark code in Scala.
+
+```
+%%sql project.redshift
+select current_user()
+```
+
+### Some other helpful magics
+
+```
+%help - list available magics and related information
+
+%send_to_remote - send local variable to remote compute
+
+%%configure - configure spark application config in remote compute
+```

sagemaker_studio_dataengineering_sessions-1.0.2/README.md

@@ -0,0 +1,98 @@
+# SageMakerStudioDataEngineeringSessions
+
+SageMaker Unified Studio Data Engineering Sessions
+
+This pacakge depends on SageMaker Unified Studio environment, if you are using SageMaker Unified Studio, see [AWS Doc](https://docs.aws.amazon.com/sagemaker-unified-studio/latest/userguide/what-is-sagemaker-unified-studio.html) for guidance.
+
+This package contains functionality to support SageMaker Unified Studio connecting to various AWS Compute including EMR/EMR Serverless/Glue/Redshift etc. 
+
+It is utilizing [ipython magics](https://ipython.readthedocs.io/en/stable/interactive/magics.html) and [AWS DataZone Connections](https://docs.aws.amazon.com/datazone/latest/APIReference/API_ListConnections.html) to achieve the following features.
+
+## Features
+
+- Connect to remote compute
+- Execute Spark code in remote compute in Python/Scala
+- Execute SQL queries in remote compute
+- Send local variables to remote compute
+
+
+## How to setup
+
+If you are using SageMaker Unifed Studio, you can skip this part, SageMaker Unifed Studio already set up the package.
+
+This package contains various Jupyter Magics to achieve its functionality.
+
+To load these magics, make sure you have iPython config file generated. If not, you could run `ipython profile create`, then a file with path `~/.ipython/profile_default/ipython_config.py` should be generated
+
+Then you will need to add the following line in the end of that config file
+
+```
+c.InteractiveShellApp.extensions.extend(['sagemaker_studio_dataengineering_sessions.sagemaker_connection_magic'])
+```
+
+Once that is finished, you could restart the ipython kernel and run `%help` to see a list of supported magics
+
+## Examples
+
+
+To connect to remote compute, a DataZone Connection is required, you could create it via [CreateConnection API](https://docs.aws.amazon.com/datazone/latest/APIReference/API_CreateConnection.html), Let's say there's an existing connection called project.spark. 
+
+### Supported Connection Type:
+
+- IAM
+- SPARK
+- REDSHIFT
+- ATHENA
+
+### Connect to remote compute and Execute Spark Code in Python
+The following example will connect to AWS Glue Interactive session and run the spark code in Glue.
+
+```
+%%pyspark project.spark
+
+import sys
+import boto3
+from awsglue.utils import getResolvedOptions
+from pyspark.context import SparkContext
+from pyspark.sql import SparkSession
+from pyspark.sql.functions import col
+
+args = getResolvedOptions(sys.argv, ["redshift_url", "redshift_iam_role", "redshift_tempdir","redshift_jdbc_iam_url"])
+print(f"{args}")
+
+sc = SparkContext.getOrCreate()
+spark = SparkSession(sc)
+
+df = spark.read.csv(f"s3://sagemaker-example-files-prod-{boto3.session.Session().region_name}/datasets/tabular/dirty-titanic/", header=True)
+df.show(5, truncate=False)
+df.printSchema()
+
+df.createOrReplaceTempView("df_sql_tempview")
+```
+
+### Execute Spark Code in Scala
+The following example will connect to AWS Glue Interactive session and run the spark code in Scala.
+
+```
+%%scalaspark project.spark
+val dfScala = spark.sql("SELECT count(0) FROM df_sql_tempview")
+dfScala.show()
+```
+
+### Execute SQL query in remote compute
+The following example will connect to AWS Glue Interactive session and run the spark code in Scala.
+
+```
+%%sql project.redshift
+select current_user()
+```
+
+### Some other helpful magics
+
+```
+%help - list available magics and related information
+
+%send_to_remote - send local variable to remote compute
+
+%%configure - configure spark application config in remote compute
+```

sagemaker_studio_dataengineering_sessions-1.0.2/pyproject.toml

@@ -0,0 +1,136 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "sagemaker-studio-dataengineering-sessions"
+version = "1.0.2"
+readme = "README.md"
+dependencies = [
+  # Put your dependencies here!
+    "boto3",
+    "pytest",
+    "ipython",
+    # database session manager
+    "amazon_sagemaker_sql_execution>=0.0.7,<1",
+    "pandas",
+    "sqlparse>=0.4.0",
+    # spark session manager
+    "sparkmagic",
+    "botocore",
+    # connection magic
+    "panel",
+    "pytest",
+    "pytest-cov"
+]
+requires-python = ">=3.11"
+
+[tool.hatch.envs.default]
+# This controls what version of Python you want to be the default
+# when running any scripts or tools to do things like debug test failures
+# or do general development.  It's lockfile is ./requirements.txt
+python = "3.11"
+dependencies = [
+    "pytest",
+#    "mypy",
+]
+
+[tool.pytest.ini_options]
+addopts = [
+    "--durations=5",
+    "--color=yes",
+]
+testpaths = [ "tests" ]
+
+[tool.coverage.run]
+source_pkgs = ["sagemaker_studio_dataengineering_sessions"]
+branch = true
+parallel = true
+
+[tool.coverage.paths]
+"sagemaker_studio_dataengineering_sessions" = ["src/sagemaker_studio_dataengineering_sessions", "**/site-packages/sagemaker_studio_dataengineering_sessions"]
+
+[tool.coverage.report]
+fail_under = 60
+exclude_lines = [
+  "no cov",
+  "if __name__ == .__main__.:",
+  "if TYPE_CHECKING:",
+]
+show_missing = true
+
+[tool.coverage.xml]
+output = "private/brazil-documentation/coverage/coverage.xml"
+
+[tool.coverage.html]
+directory = "private/brazil-documentation/coverage/"
+
+[tool.ruff.lint]
+isort.known-first-party = ["sagemaker_studio_dataengineering_sessions"]
+exclude = [ "./build", ".hatch", "private" ]
+
+[tool.ruff.lint.per-file-ignores]
+# Tests can use magic values, assertions, and relative imports
+"tests/**/*" = ["PLR2004", "S101", "TID252"]
+
+[tool.hatch.build]
+directory = "./external-distribution"
+packages = ["src/sagemaker_studio_dataengineering_sessions"]
+
+[tool.hatch.build.targets.sdist]
+ignore-vcs = true
+
+[tool.hatch.build.targets.wheel]
+ignore-vcs = true
+
+[tool.hatch.build.hooks.custom]
+path = "scripts/build_hooks.py"
+
+[tool.hatch.env]
+requires = [ "hatch-pip-compile" ]
+
+[tool.hatch.envs.default.scripts]
+# These are scripts you can run using `brazil-build run <script-name>`
+#typing = [
+#  "mkdir -p .mypy_cache",
+#  "mypy --install-types --non-interactive src/sagemaker_studio_dataengineering_sessions tests"
+#]
+
+# This command is for updating all your lock files across all environments 
+update = [ "hatch-pip-compile --upgrade --all" ]
+
+release = [
+#    "typing",
+#    "hatch test --all --cover",
+    "pytest tests --cov"
+]
+
+[[tool.hatch.envs.hatch-test.matrix]]
+# This defines multiple variables you can generate combinations
+# to test underneath different environments.  A separate environment and 
+# lock file will be created for every combination located in `./requirements/`
+python = ["3.11", "3.12"]
+
+## This environment is used solely to generate a lock file on hatch,
+# and hatch-pip-compile that can be automatically updated
+[tool.hatch.envs.build-tools]
+# This version states what version your build tools build with.  To change it,
+# you will need to:
+# * Remove the `requirements/requirements-build-tools.txt` file
+# * Run `brazil-build run update` to generate a new lock file for the environment
+python = "3.11"
+detached = true
+skip-install = true
+dependencies = [
+  "hatch",
+  "hatch-pip-compile",
+]
+
+# PeruHatch repository and package locking plugin
+[tool.hatch.env.collectors.custom]
+path = ".hatch/hatch_plugin.py"
+
+# This is necessary to use 'uv' as the resolver if this is the top-level package
+# in a monorepo (which is usually the case).  Remove this if copying the
+# package into a monorepo
+[tool.uv.workspace]

sagemaker_studio_dataengineering_sessions-1.0.2/sagemaker_studio_dataengineering_sessions/sagemaker_base_session_manager/__init__.py

@@ -0,0 +1,7 @@
+import logging
+
+from sagemaker_studio_dataengineering_sessions.sagemaker_base_session_manager.common.logger_utils import setup_logger
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+setup_logger(logger, "SageMakerBaseSessionManager", "connection_magic")

sagemaker_studio_dataengineering_sessions-1.0.2/sagemaker_studio_dataengineering_sessions/sagemaker_base_session_manager/base_session_manager.py

@@ -0,0 +1,162 @@
+import abc
+import logging
+
+from sagemaker_studio_dataengineering_sessions.sagemaker_base_session_manager.common.constants import Language
+from sagemaker_studio_dataengineering_sessions.sagemaker_base_session_manager.common.sagemaker_connection_display import SageMakerConnectionDisplay
+
+
+class BaseSessionManager(metaclass=abc.ABCMeta):
+
+    def __init__(self) -> None:
+        # The limit for the number of rows that can be returned in a SQL query. Defaults to 10000.
+        self.sql_result_row_limit = 10_000
+
+    @abc.abstractmethod
+    def create_session(self):
+        """
+        Create a new session
+        """
+        raise NotImplementedError('Must define create_session to use this BaseSessionManager')
+
+    @abc.abstractmethod
+    def run_statement(self, cell="", language: Language = None):
+        """
+        Run a statement against the session
+        :param cell: the code to run
+        """
+        raise NotImplementedError('Must define run_statement to use this BaseSessionManager')
+
+    @abc.abstractmethod
+    def stop_session(self):
+        """
+        Stop the session
+        """
+        raise NotImplementedError('Must define stop_session to use this BaseSessionManager')
+
+    @abc.abstractmethod
+    def is_session_connectable(self) -> bool:
+        """
+        Check if the session is connectable and in a valid status for running statement
+        :return: true if the session is in a valid status, false otherwise
+        """
+        raise NotImplementedError('Must define is_session_connectable to use this BaseSessionManager')
+
+    @abc.abstractmethod
+    def _configure_core(self, cell: str):
+        """
+        To configure current compute to be connected and start a new session with applied configuration
+        :param cell: the content to be applied for the session
+        """
+
+    def configure(self, cell: str, force: bool = False):
+        """
+        To configure current compute to be connected and start a new session with applied configuration
+        :param cell: the content to be applied for the session
+        :param force: a boolean to check if a user wants to force start a new session to apply configuration
+        """
+        if self.is_session_connectable():
+            if not force:
+                SageMakerConnectionDisplay.send_error(
+                    "A session has already been started. If you intend to recreate the "
+                    "session with new configurations, please include the -f or --force argument.")
+            else:
+                self.stop_session()
+                self._configure_core(cell)
+                self.create_session()
+        else:
+            self._configure_core(cell)
+        return
+
+    def send_to_remote(self, local_var: str, remote_var: str, language=Language.python):
+        # Not an abstract method because by default a session manager does not support this
+        # only Spark session manager supports this for now.
+        """
+        Send a local variable in kernel's userspace to remote compute.
+        e.g: for an EMR cluster, send a local variable to spark
+        :param local_var: local variable name
+        :param remote_var: remote variable name
+        """
+        raise NotImplementedError('Send_to_remote is not supported')
+
+    def get_info(self):
+        """
+        Get information about the connected compute session
+        """
+        raise NotImplementedError('get_info is not supported for current session')
+
+    def get_session_id(self):
+        """
+        Get the session ID of the connected compute session
+        """
+        raise NotImplementedError('get_session_id is not supported for current session')
+
+    def get_status(self):
+        """
+        Get the status of the connected compute session
+        """
+        raise NotImplementedError('get_status is not supported for current session')
+
+    def add_tags(self, tags: str):
+        """
+        Add tags to the connected compute session resources
+        """
+        raise NotImplementedError('add_tags is not supported for current session')
+
+    def get_logs(self):
+        """
+        Gets the current session's Livy logs
+        """
+        raise NotImplementedError('get_logs is not supported for current session')
+
+    def matplot(self, line: str):
+        """
+        Using matplotlib to plot the current session's plot'
+        """
+        raise NotImplementedError('matplot is not supported for current session')
+
+    def set_session_id_prefix(self, prefix: str, force: bool = False):
+        """
+        Sets the session ID prefix of the session to be created
+        """
+        raise NotImplementedError('set_session_id_prefix is not supported for current session')
+
+    def set_number_of_workers(self, number: str, force: bool = False):
+        """
+        Sets the number of workers of the session to be created
+        """
+        raise NotImplementedError('set_number_of_workers is not supported for current session')
+
+    def set_worker_type(self, type: str, force: bool = False):
+        """
+        Sets the worker type of the session to be created
+        """
+        raise NotImplementedError('set_worker_type is not supported for current session')
+
+    def set_session_type(self, session_type: str, force: bool = False):
+        """
+        Sets the session type of the session to be created. Acceptable session_type values are: streaming and etl.
+        """
+        raise NotImplementedError('set_session_type is not supported for current session')
+
+    def set_glue_version(self, glue_version: str, force: bool = False):
+        """
+        Sets the glue version of the session to be created
+        """
+        raise NotImplementedError('set_glue_version is not supported for current session')
+
+    def set_idle_timeout(self, idle_timeout: str, force: bool = False):
+        """
+        Sets the idle timeout value of the session to be created
+        """
+        raise NotImplementedError('set_idle_timeout is not supported for current session')
+
+    def spark_conf(self, spark_conf: str, force: bool = False):
+        """
+        Sets the spark configuration value of the session to be created
+        """
+        raise NotImplementedError('spark_conf is not supported for current session')
+
+    def get_logger(self):
+        if not hasattr(self, '_logger'):
+            self._logger = logging.getLogger(__name__)
+        return self._logger