feature: add support for Spark Connect

.github/workflows/test.yml

@@ -6,6 +6,7 @@ on:
   pull_request:
     branches:
     - main
+    - release/*
   workflow_dispatch:
     inputs:
       logLevel:
@@ -40,8 +41,8 @@ jobs:
         fetch-depth: 0
         ref: ${{ github.event.pull_request.head.ref }}
         repository: ${{ github.event.pull_request.head.repo.full_name }}
-    - name: Fetch main branch
-      run: git fetch origin main:main
+    - name: Fetch target branch
+      run: git fetch origin ${{ github.event.pull_request.base.ref || 'main'}}:${{ github.event.pull_request.base.ref || 'main'}}
     - name: Check changes
       id: check
       run: |
@@ -71,10 +72,12 @@ jobs:
 #        os: [ubuntu-latest, windows-latest, macos-latest]  # FIXME: Add Windows and macOS
         os: [ubuntu-latest]
         python-version: ['3.9', '3.10', '3.11', '3.12']
-        pyspark-version: ['33', '34', '35']
+        pyspark-version: ['33', '34', '35', '35r']
         exclude:
           - python-version: '3.9'
             pyspark-version: '35'
+          - python-version: '3.9'
+            pyspark-version: '35r'
           - python-version: '3.11'
             pyspark-version: '33'
           - python-version: '3.11'
@@ -100,7 +103,7 @@ jobs:
 #        hatch fmt --check --python=${{ matrix.python-version }}
 
     - name: Run tests
-      run: hatch test --python=${{ matrix.python-version }} -i version=pyspark${{ matrix.pyspark-version }}
+      run: hatch test --python=${{ matrix.python-version }} -i version=pyspark${{ matrix.pyspark-version }} --verbose
 
   # https://github.com/marketplace/actions/alls-green#why
   final_check: # This job does nothing and is only used for the branch protection

makefile

@@ -105,16 +105,16 @@ coverage: cov
 all-tests:
 	@echo "\033[1mRunning all tests:\033[0m\n\033[35m This will run the full test suite\033[0m"
 	@echo "\033[1;31mWARNING:\033[0;33m This may take upward of 20-30 minutes to complete!\033[0m"
-	@hatch test --no-header --no-summary
+	@hatch test --no-header
 .PHONY: spark-tests  ## testing - Run SPARK tests in ALL environments
 spark-tests:
 	@echo "\033[1mRunning Spark tests:\033[0m\n\033[35m This will run the Spark test suite against all specified environments\033[0m"
 	@echo "\033[1;31mWARNING:\033[0;33m This may take upward of 20-30 minutes to complete!\033[0m"
-	@hatch test -m spark --no-header --no-summary
+	@hatch test -m spark --no-header
 .PHONY: non-spark-tests  ## testing - Run non-spark tests in ALL environments
 non-spark-tests:
 	@echo "\033[1mRunning non-Spark tests:\033[0m\n\033[35m This will run the non-Spark test suite against all specified environments\033[0m"
-	@hatch test -m "not spark" --no-header --no-summary
+	@hatch test -m "not spark" --no-header
 
 .PHONY: dev-test ## testing - Run pytest, with all tests in the dev environment
 dev-test:

pyproject.toml

@@ -57,19 +57,30 @@ async_http = [
   "nest-asyncio>=1.6.0",
 ]
 box = ["boxsdk[jwt]==3.8.1"]
-pandas = ["pandas>=1.3", "setuptools", "numpy<2.0.0"]
+pandas = ["pandas>=1.3", "setuptools", "numpy<2.0.0", "pandas-stubs"]
 pyspark = ["pyspark>=3.2.0", "pyarrow>13"]
+pyspark_connect = ["pyspark[connect]>=3.5"]
 se = ["spark-expectations>=2.1.0"]
 # SFTP dependencies in to_csv line_iterator
 sftp = ["paramiko>=2.6.0"]
 delta = ["delta-spark>=2.2"]
 excel = ["openpyxl>=3.0.0"]
 # Tableau dependencies
-tableau = [
-  "tableauhyperapi>=0.0.19484",
-  "tableauserverclient>=0.25",
+tableau = ["tableauhyperapi>=0.0.19484", "tableauserverclient>=0.25"]
+# Snowflake dependencies
+snowflake = ["snowflake-connector-python>=3.12.0"]
+# Development dependencies
+dev = [
+  "black",
+  "isort",
+  "ruff",
+  "mypy",
+  "pylint",
+  "colorama",
+  "types-PyYAML",
+  "types-requests",
+
 ]
-dev = ["black", "isort", "ruff", "mypy", "pylint", "colorama", "types-PyYAML"]
 test = [
   "chispa",
   "coverage[toml]",
@@ -186,6 +197,7 @@ features = [
   "excel",
   "se",
   "box",
+  "snowflake",
   "tableau",
   "dev",
 ]
@@ -216,6 +228,7 @@ lint = ["- ruff-fmt", "- mypy-check", "pylint-check"]
 log-versions = "python --version && {env:HATCH_UV} pip freeze | grep pyspark"
 test = "- pytest{env:HATCH_TEST_ARGS:} {args} -n 2"
 spark-tests = "test -m spark"
+spark-remote-tests = "test -m spark -m \"not skip_on_remote_session\""
 non-spark-tests = "test -m \"not spark\""
 
 # scripts.run = "echo bla {env:HATCH_TEST_ARGS:} {args}"
@@ -251,17 +264,18 @@ features = [
   "sftp",
   "delta",
   "excel",
+  "snowflake",
   "tableau",
   "dev",
   "test",
 ]
 
 parallel = true
 retries = 2
-retry-delay = 1
+retry-delay = 3
 
 [tool.hatch.envs.hatch-test.scripts]
-run = "pytest{env:HATCH_TEST_ARGS:} {args} -n auto"
+run = "pytest{env:HATCH_TEST_ARGS:} {args}"
 run-cov = "coverage run -m pytest{env:HATCH_TEST_ARGS:} {args}"
 cov-combine = "coverage combine"
 cov-report = "coverage report"
@@ -273,11 +287,11 @@ version = ["pyspark33", "pyspark34"]
 
 [[tool.hatch.envs.hatch-test.matrix]]
 python = ["3.10"]
-version = ["pyspark33", "pyspark34", "pyspark35"]
+version = ["pyspark33", "pyspark34", "pyspark35", "pyspark35r"]
 
 [[tool.hatch.envs.hatch-test.matrix]]
 python = ["3.11", "3.12"]
-version = ["pyspark35"]
+version = ["pyspark35", "pyspark35r"]
 
 [tool.hatch.envs.hatch-test.overrides]
 matrix.version.extra-dependencies = [
@@ -299,6 +313,9 @@ matrix.version.extra-dependencies = [
   { value = "pyspark>=3.5,<3.6", if = [
     "pyspark35",
   ] },
+  { value = "pyspark[connect]>=3.5,<3.6", if = [
+    "pyspark35r",
+  ] },
 ]
 
 name.".*".env-vars = [
@@ -308,10 +325,19 @@ name.".*".env-vars = [
   { key = "KOHEESIO__PRINT_LOGO", value = "False" },
 ]
 
+name.".*(pyspark35r).*".env-vars = [
+  # enable soark connect, setting to local as it will trigger 
+  # spark to start local spark server and enbale remote session
+  { key = "SPARK_REMOTE", value = "local" },
+  { key = "SPARK_TESTING", value = "True" },
+]
+
+
 [tool.pytest.ini_options]
 addopts = "-q --color=yes --order-scope=module"
 log_level = "CRITICAL"
 testpaths = ["tests"]
+asyncio_default_fixture_loop_scope = "scope"
 markers = [
   "default: added to all tests by default if no other marker expect of standard pytest markers is present",
   "spark: mark a test as a Spark test",
@@ -325,10 +351,18 @@ filterwarnings = [
   # pyspark.pandas warnings
   "ignore:distutils.*:DeprecationWarning:pyspark.pandas.*",
   "ignore:'PYARROW_IGNORE_TIMEZONE'.*:UserWarning:pyspark.pandas.*",
+  # pydantic warnings
+  "ignore:A custom validator is returning a value other than `self`.*.*:UserWarning:pydantic.main.*",
+  # pyspark.sql.connect warnings
+  "ignore:is_datetime64tz_dtype.*:DeprecationWarning:pyspark.sql.connect.*",
+  "ignore:distutils.*:DeprecationWarning:pyspark.sql.connect.*",
+  # pyspark.sql.pandas warnings
   "ignore:distutils.*:DeprecationWarning:pyspark.sql.pandas.*",
   "ignore:is_datetime64tz_dtype.*:DeprecationWarning:pyspark.sql.pandas.*",
+  "ignore:is_categorical_dtype.*:DeprecationWarning:pyspark.sql.pandas.*",
+  "ignore:iteritems.*:FutureWarning:pyspark.sql.pandas.*",
   # Koheesio warnings
-  "ignore:DayTimeIntervalType .*:UserWarning:koheesio.spark.snowflake.*",
+  "ignore:DayTimeIntervalType.*:UserWarning:koheesio.spark.snowflake.*",
 ]
 
 [tool.coverage.run]
@@ -403,16 +437,16 @@ features = [
   "box",
   "pandas",
   "pyspark",
-  "se",
+  #   "se",
   "sftp",
+  "snowflake",
   "delta",
   "excel",
   "tableau",
   "dev",
   "test",
   "docs",
 ]
-extra-dependencies = ["pyspark==3.4.*"]
 
 
 ###  ~~~~~~~~~~~~~~~~~~  ###
@@ -569,10 +603,18 @@ unfixable = []
 dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
 
 [tool.mypy]
+python_version = "3.10"
+files = ["koheesio/**/*.py"]
+plugins = ["pydantic.mypy"]
+pretty = true
+warn_unused_configs = true
 check_untyped_defs = false
 disallow_untyped_calls = false
 disallow_untyped_defs = true
-files = ["koheesio/**/*.py"]
+warn_no_return = false
+implicit_optional = true
+allow_untyped_globals = true
+disable_error_code = ["attr-defined", "return-value", "union-attr", "override"]
 
 [tool.pylint.main]
 fail-under = 9.5

src/koheesio/__about__.py

@@ -12,7 +12,7 @@
 
 LICENSE_INFO = "Licensed as Apache 2.0"
 SOURCE = "https://github.com/Nike-Inc/koheesio"
-__version__ = "0.8.1"
+__version__ = "0.9.0rc0"
 __logo__ = (
     75,
     (
@@ -32,7 +32,7 @@
 
 
 # fmt: off
-def _about():  # pragma: no cover
+def _about() -> str:  # pragma: no cover
     """Return the Koheesio logo and version/about information as a string
     Note: this code is not meant to be readable, instead it is written to be as compact as possible
     """

src/koheesio/__init__.py

@@ -25,7 +25,7 @@
 ]
 
 
-def print_logo():
+def print_logo() -> None:
     global _logo_printed
     global _koheesio_print_logo
 

src/koheesio/asyncio/__init__.py

@@ -16,15 +16,13 @@ class AsyncStepMetaClass(StepMetaClass):
     It inherits from the StepMetaClass and provides additional functionality for
     executing asynchronous steps.
 
-    Attributes:
-        None
-
-    Methods:
-        _execute_wrapper: Wrapper method for executing asynchronous steps.
+    Methods
+    -------
+    _execute_wrapper: Wrapper method for executing asynchronous steps.
 
     """
 
-    def _execute_wrapper(cls, *args, **kwargs):
+    def _execute_wrapper(cls, *args, **kwargs):  # type: ignore[no-untyped-def]
         """Wrapper method for executing asynchronous steps.
 
         This method is called when an asynchronous step is executed. It wraps the
@@ -60,16 +58,14 @@ class AsyncStepOutput(Step.Output):
         Merge key-value map with self.
     """
 
-    def merge(self, other: Union[Dict, StepOutput]):
+    def merge(self, other: Union[Dict, StepOutput]) -> "AsyncStepOutput":
         """Merge key,value map with self
 
         Examples
         --------
         ```python
         step_output = StepOutput(foo="bar")
-        step_output.merge(
-            {"lorem": "ipsum"}
-        )  # step_output will now contain {'foo': 'bar', 'lorem': 'ipsum'}
+        step_output.merge({"lorem": "ipsum"})  # step_output will now contain {'foo': 'bar', 'lorem': 'ipsum'}
         ```
 
         Functionally similar to adding two dicts together; like running `{**dict_a, **dict_b}`.
@@ -89,6 +85,7 @@ def merge(self, other: Union[Dict, StepOutput]):
         return self
 
 
+# noinspection PyUnresolvedReferences
 class AsyncStep(Step, ABC, metaclass=AsyncStepMetaClass):
     """
     Asynchronous step class that inherits from Step and uses the AsyncStepMetaClass metaclass.