move src code to package pipelines

Author: Joshua Cook

.gitignore

@@ -6,3 +6,6 @@
 **/.local
 **/.pytest_cache
 **/__pycache__
+**/tests/data/bronze
+**/tests/data/silver
+**/tests/data/checkpoints

env/docker/base/Dockerfile

@@ -1,3 +1,4 @@
 FROM jupyter/pyspark-notebook
 COPY requirements.txt .
 RUN pip install -r requirements.txt
+ENV PYTHONPATH ${PYTHONPATH}:/home/jovyan/src

example/databricks_job_with_wheel/env/docker/base/Dockerfile

@@ -1,3 +1,4 @@
 FROM jupyter/pyspark-notebook
 COPY requirements.txt .
 RUN pip install -r requirements.txt
+ENV PYTHONPATH ${PYTHONPATH}:/home/jovyan/src

example/databricks_job_with_wheel/src/__init__.py renamed to example/databricks_job_with_wheel/src/pipelines/__init__.py

@@ -1,11 +1,12 @@
 import os
 import pytest
+from shutil import rmtree
 from pyspark.sql import DataFrame, SparkSession
 from pyspark.sql.types import StructType
 
-from src.config import generate_paths, generate_schemas
-from src.operations import create_stream_writer, transform_bronze, transform_raw
-from src.utility import generate_spark_session, read_stream_json
+from pipelines.config import generate_paths, generate_schemas
+from pipelines.operations import create_stream_writer, transform_bronze, transform_raw
+from pipelines.utility import generate_spark_session, load_delta_table, read_stream_json
 
 
 @pytest.fixture(scope="module")
@@ -20,17 +21,23 @@ def spark() -> SparkSession:
 
 @pytest.fixture(scope="module")
 def raw_schema() -> StructType:
-    yield generate_schemas("RAW_SCHEMA")
+    yield generate_schemas("raw_schema")
+
+
+@pytest.fixture(scope="module")
+def bronze_path(env: str) -> StructType:
+    path = generate_paths(env, "bronze")
+    yield path
 
 
 @pytest.fixture(scope="module")
 def bronze_schema() -> StructType:
-    yield generate_schemas("BRONZE_SCHEMA")
+    yield generate_schemas("bronze_schema")
 
 
 @pytest.fixture(scope="module")
 def silver_schema() -> StructType:
-    yield generate_schemas("SILVER_SCHEMA")
+    yield generate_schemas("silver_schema")
 
 
 @pytest.fixture()
@@ -45,12 +52,28 @@ def bronze_df(spark: SparkSession, env: str, bronze_schema: StructType) -> DataF
     yield read_stream_json(spark, bronze_path, bronze_schema)
 
 
+@pytest.fixture()
+def silver_df(spark: SparkSession, env: str, silver_schema: StructType) -> DataFrame:
+    test_silver_path = generate_paths(env, "test_silver")
+    silver_path = generate_paths(env, "silver")
+    silver_json_df = read_stream_json(spark, test_silver_path, silver_schema)
+    (
+        silver_json_df.writeStream.format("delta")
+        .partitionBy("p_eventdate")
+        .option("path", silver_path)
+        .save()
+    )
+    yield load_delta_table(spark, silver_path)
+    rmtree(silver_path)
+
+
 class TestSparkDataframeOperations:
-    def test_create_stream_write(self, env, spark, raw_df, bronze_schema):
+    def test_create_stream_write(self, env, spark, raw_df, bronze_path, bronze_schema):
         transformed_raw_df = transform_raw(spark, raw_df)
         bronze_checkpoint = generate_paths(env, "bronze_checkpoint")
         raw_to_bronze_writer = create_stream_writer(
             dataframe=transformed_raw_df,
+            path=bronze_path,
             checkpoint=bronze_checkpoint,
             name="write_raw_to_bronze",
             partition_column="p_ingestdate",
@@ -64,3 +87,6 @@ def test_transform_raw(self, spark, raw_df, bronze_schema):
     def test_transform_bronze(self, spark, bronze_df, silver_schema):
         transformed_bronze_df = transform_bronze(spark, bronze_df)
         assert transformed_bronze_df.schema == silver_schema
+
+    def test_prepare_interpolation_dataframe(self, spark, silver_df):
+        pass

example/databricks_job_with_wheel/src/config.py renamed to example/databricks_job_with_wheel/src/pipelines/config.py

@@ -4,9 +4,9 @@
 from pyspark.sql import DataFrame, SparkSession
 from pyspark.sql.types import StructType
 
-from src.config import generate_paths, generate_schemas
-from src.operations import create_stream_writer, transform_raw
-from src.utility import (
+from pipelines.config import generate_paths, generate_schemas
+from pipelines.operations import create_stream_writer, transform_raw, transform_bronze
+from pipelines.utility import (
     generate_spark_session,
     load_delta_table,
     read_stream_json,
@@ -24,6 +24,12 @@ def spark() -> SparkSession:
     yield generate_spark_session()
 
 
+@pytest.fixture()
+def bronze_df(spark: SparkSession, env: str, bronze_schema: StructType) -> DataFrame:
+    bronze_path = generate_paths(env, "test_bronze")
+    yield read_stream_json(spark, bronze_path, bronze_schema)
+
+
 @pytest.fixture(scope="module")
 def bronze_checkpoint(env: str) -> StructType:
     path = generate_paths(env, "bronze_checkpoint")
@@ -40,13 +46,37 @@ def bronze_path(env: str) -> StructType:
 
 @pytest.fixture(scope="module")
 def bronze_schema() -> StructType:
-    yield generate_schemas("BRONZE_SCHEMA")
+    yield generate_schemas("bronze_schema")
+
+
+@pytest.fixture(scope="module")
+def raw_schema() -> StructType:
+    yield generate_schemas("raw_schema")
 
 
 @pytest.fixture()
-def raw_df(spark: SparkSession, env: str) -> DataFrame:
+def raw_df(spark: SparkSession, env: str, raw_schema: StructType) -> DataFrame:
     raw_path = generate_paths(env, "test_raw")
-    yield read_stream_json(spark, raw_path)
+    yield read_stream_json(spark, raw_path, raw_schema)
+
+
+@pytest.fixture(scope="module")
+def silver_checkpoint(env: str) -> StructType:
+    path = generate_paths(env, "silver_checkpoint")
+    yield path
+    rmtree(path)
+
+
+@pytest.fixture(scope="module")
+def silver_path(env: str) -> StructType:
+    path = generate_paths(env, "silver")
+    yield path
+    rmtree(path)
+
+
+@pytest.fixture(scope="module")
+def silver_schema() -> StructType:
+    yield generate_schemas("silver_schema")
 
 
 class TestSparkIntegrations:
@@ -57,11 +87,29 @@ def test_raw_to_bronze(
         transformed_raw_df = transform_raw(spark, raw_df)
         raw_to_bronze_writer = create_stream_writer(
             dataframe=transformed_raw_df,
+            path=bronze_path,
             checkpoint=bronze_checkpoint,
             name=stream_name,
             partition_column="p_ingestdate",
         )
-        raw_to_bronze_writer.start(bronze_path)
+        raw_to_bronze_writer.start()
 
         until_stream_is_ready(spark, stream_name)
         assert load_delta_table(spark, bronze_path).count() == 7320
+
+    def test_bronze_to_silver(
+        self, env, spark, bronze_df, silver_checkpoint, silver_path, silver_schema
+    ):
+        stream_name = "write_bronze_to_silver"
+        transformed_bronze_df = transform_bronze(spark, bronze_df)
+        bronze_to_silver_writer = create_stream_writer(
+            dataframe=transformed_bronze_df,
+            path=silver_path,
+            checkpoint=silver_checkpoint,
+            name=stream_name,
+            partition_column="p_eventdate",
+        )
+        bronze_to_silver_writer.start()
+
+        until_stream_is_ready(spark, stream_name)
+        assert load_delta_table(spark, silver_path).count() == 7320