fix running main

Author: fjakobs

default_python/scratch/exploration.ipynb

@@ -1,5 +1,15 @@
 {
  "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -21,7 +31,7 @@
     "sys.path.append('../src')\n",
     "from default_python import main\n",
     "\n",
-    "main.get_taxis().show(10)"
+    "main.get_taxis(spark).show(10)"
    ]
   }
  ],
@@ -41,8 +51,16 @@
    "name": "python3"
   },
   "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
    "name": "python",
-   "version": "3.11.4"
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.8"
   }
  },
  "nbformat": 4,

default_python/src/default_python/main.py

@@ -1,9 +1,11 @@
 from pyspark.sql import SparkSession
 
-def get_taxis(spark: SparkSession):
+def get_taxis(spark: SparkSession):    
   return spark.read.table("samples.nyctaxi.trips")
 
 def main():
+  from databricks.connect import DatabricksSession as SparkSession
+  spark = SparkSession.builder.getOrCreate()
   get_taxis(spark).show(5)
 
 if __name__ == '__main__':

default_python/src/notebook.ipynb

@@ -75,7 +75,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.6"
+   "version": "3.11.8"
   }
  },
  "nbformat": 4,

default_python/tests/main_test.py

@@ -1,15 +1,15 @@
 from databricks.connect import DatabricksSession as SparkSession
+from pytest import fixture
 from default_python import main
 from pytest import fixture
 
-# Create a new Databricks Connect session. If this fails,
-# check that you have configured Databricks Connect correctly.
-# See https://docs.databricks.com/dev-tools/databricks-connect.html.
-
 @fixture(scope="session")
 def spark():
-    return SparkSession.builder.getOrCreate()
+    spark = SparkSession.builder.getOrCreate()
+    yield spark
+    spark.stop()
+
 
-def test_main(spark):
+def test_main(spark: SparkSession):
     taxis = main.get_taxis(spark)
     assert taxis.count() > 5