Inspiration from Batch Processing

Batch Processing shines with below (though not limited) general use cases

• Transforming data in ETL/ELT data pipelines.

• Performing aggregation, grouping, filtering, joins, analytics, etc, the list is never ending.

• Doing all kinds of operations on table data like append, merge, delete, update, etc.

This side of the world is pretty mature now. We have a very good set of tools, frameworks, etc. that allows us to develop the pipeline. So why not apply the same functionality on a streaming pipeline but at the same time not lose its real-time processing character.

Problem Statement

Track customers' journeys on an e-commerce website. The events include which products a customer viewed, which products were added to the cart, and finally, which products were purchased. Calculate metrics that can be used for analysis.

I am using Kafka as my streaming platform. In Python, two popular client libraries can be used to interact with Kafka brokers

1. confluent kafka

2. kafka python

Here is my alternative take - neither of them should be used directly.

• Neither of them has Python’s Typed support, so no IDE can work efficiently.

• Both need a fair amount of boilerplate code

• Their scope of work is limited to just Producer, Consumer & Admin.

• To apply principles of Batch processing, we have to write a lot of custom code.

You need not worry, there are open-source libraries like Pathway and Quix Streams that can be used for stream processing. Let’s quickly compare them

• Pathway has many more GitHub stars compared to Quix Streams.

• Pathway also supports a lot more connectors and sinks than the other.

• Both allow us to do batch processing tasks like filters, aggregation, joins, analytics, and Python UDF on streaming data.

• Pathway's strength in supporting many connectors and sinks is also its weakness, as it makes the library quite large. Quix Streams is designed for one main purpose: Streaming.

• The biggest advantage of Quix Streams is that it lets us use low-level Kafka client tools with full Typed support, allowing us to use IDEs efficiently. This was the main reason I chose Quix Streams.

The Producer:

Below, I am generating Fake data for a Customer journey & sending the events into their respective topics.

 import logging

from quixstreams import Application

from src.ops.generator import generate_dummy_e_commerce_data

logger = logging.getLogger("data-pipeline")
app = Application(
    broker_address="localhost:9092",
    loglevel="DEBUG",
    producer_extra_config={
        "linger.ms": "300",  # Wait up to 300ms for more messages before sending
        "compression.type": "gzip",  # Use gzip compression for messages
    },
)
product_view_topic = app.topic(
    "ecomm_product_view", value_serializer="json", key_serializer="string"
)
# Creating the kafka Topic client
cart_topic = app.topic("ecomm_cart", value_serializer="json", key_serializer="string")
buy_topic = app.topic("ecomm_buy", value_serializer="json", key_serializer="string")

# quixstreams' Application allow us to use context manager so we won't need to worry 
# about closing connection gracefully
with app.get_producer() as producer:
    for _ in range(200):
        logger.info(f"current iteration: {_}")
        # generate dummy data
        event = generate_dummy_e_commerce_data()

        # sending the event to respective topics
        # Product view event
        product_view_msg = product_view_topic.serialize(
            key=event["product_view"].user_id,
            value=event["product_view"].model_dump(mode="json"),
        )
        producer.produce(
            product_view_topic.name,
            value=product_view_msg.value,
            key=product_view_msg.key,
        )
        logger.debug(f"producing product view event: {event['product_view'].event_id}")

        # Add to cart event
        # NOTE - It is possible that the add_to_cart and purchase events are None
        if event["add_to_cart"]:
            cart_msg = cart_topic.serialize(
                key=event["add_to_cart"].user_id,
                value=event["add_to_cart"].model_dump(mode="json"),
            )
            producer.produce(cart_topic.name, value=cart_msg.value, key=cart_msg.key)
            logger.debug(
                f"producing product add to cart event: {event['add_to_cart'].event_id}"
            )

        if event["purchase"]:
            buy_msg = buy_topic.serialize(
                key=event["purchase"].user_id,
                value=event["purchase"].model_dump(mode="json"),
            )
            producer.produce(buy_topic.name, value=buy_msg.value, key=buy_msg.key)
            logger.debug(f"producing buy event: {event['purchase'].event_id}")

Important Point:

• The default settings for Producer are good, but to extract more performance, you’ll need to fine-tune them. I usually play around following the producer config setting:

  • compression.type: To compress messages. I usually use gzip. It has wide support.

  • linger.ms: Duration to wait for more gathering messages before sending. It should depend on the frequency of incoming messages.

• It’s recommended to use the context manager of quixstreams app object, so that towards the end connections will be closed gracefully.

The Consumer: Part A

Consuming messages from three topics —> Perform data processing —> Perform join to calculate customer journey —> Publish results to 4th topic.

from quixstreams import Application

from src.ops.transform import convert_utc_to_ist

app = Application(
    broker_address="localhost:9092",
    consumer_group="ecomm_sync_group",
    auto_offset_reset="earliest",
    consumer_extra_config={
        "auto.offset.reset": "earliest",  # Start reading from the earliest message
        "enable.auto.commit": "true",  # Automatically commit offsets
    },
    loglevel="DEBUG",
)

# add topics to consume
product_view_topic = app.topic(
    "ecomm_product_view", value_serializer="json", key_serializer="string"
)
cart_topic = app.topic("ecomm_cart", value_deserializer="json", key_serializer="string")
buy_topic = app.topic("ecomm_buy", value_deserializer="json", key_serializer="string")
# Output topic for customer journey
customer_journey_topic = app.topic(
    "ecomm_customer_journey",
    value_serializer="json",
)

# Streaming dataframe consumers
product_view_sdf = app.dataframe(product_view_topic)
cart_sdf = app.dataframe(cart_topic)
buy_sdf = app.dataframe(buy_topic)

# Stream processing
# Product view topic
product_view_sdf["timestamp"] = product_view_sdf["timestamp"].apply(convert_utc_to_ist)

# Add to cart topic
cart_sdf["price"] = cart_sdf["price"].apply(lambda x: round(x, 2))
cart_sdf["timestamp"] = cart_sdf["timestamp"].apply(convert_utc_to_ist)
# Join product view and cart dataframes on user_id
joined_view_cart = cart_sdf.join_asof(
    right=product_view_sdf, how="left", on_merge="keep-left"
)

# Buy product topic
buy_sdf["price"] = buy_sdf["price"].apply(lambda x: round(x, 2))
buy_sdf["timestamp"] = buy_sdf["timestamp"].apply(convert_utc_to_ist)
joined_buy = buy_sdf.join_asof(right=joined_view_cart, how="left", on_merge="keep-left")
joined_buy.to_topic(customer_journey_topic)

# Starting the app to process streams real-time
app.run()

Important Points:

• Same as Producer, the default settings for Consumer are good, but to extract more performance, you’ll need to fine-tune them. I usually play around following the producer config setting:

• Quix streams provides us with StreamingDataFrame interface to apply all kinds of transformation & analytical logic that we want to apply, including Python UDF.

• I am performing a Stateful Join on two StreamingDataFrame. It uses RocksDB to maintain the State with flushing data. I will highly recommend reading more on it here

• After a couple of Join Operation, I am sending data as an event to yet another topic. This will only get triggered when the join operation finds a key to join. In this way, I can track the entire journey of the customer from product view to add to cart to finally buy.

The Consumer: Part B

Consuming messages using a low-level client library API to perform custom sink logic.

import logging
from pathlib import Path

import orjson
import polars as pl
from quixstreams import Application

from src.connector import DataWriter

logger = logging.getLogger("data-pipeline")
merge_option = DataWriter.generate_delta_table_merge_method_options(
    when_not_matched_insert_all=True, when_matched_update_all=True
)
app = Application(
    broker_address="localhost:9092",
    consumer_group="ecomm_customer_report_group",
    auto_offset_reset="earliest",
    consumer_extra_config={
        "enable.auto.commit": True,  # Automatically commit offsets
    },
    loglevel="DEBUG",
)

with app.get_consumer() as consumer:
    consumer.subscribe(topics=["ecomm_customer_journey"])

    # Starting the 'Forever consuming consumer'
    while True:
        message = consumer.poll(0.5)
        if message is None:
            continue
        elif message.error():
            logger.error("Kafka error:", message.error())
            continue

        value = message.value()

        # Merge data into Delta lake table
        df = pl.from_dict(orjson.loads(value))
        merge_stats = DataWriter.delta_table_merge_disk(
            df=df,
            path=Path(__file__).parent.parent.parent
            / "data/gold/ecomm_customer_journey",
            delta_merge_options={
                "predicate": "source.event_id = target.event_id",  # condition to determine upsert req
                "source_alias": "source",
                "target_alias": "target",
            },
            delta_merge_method_options=merge_option,
        )
        logger.info(f"merge successfully with stats: {merge_stats}")

        consumer.store_offsets(message=message)

Important Points:

• This is the second way to get messages from Kafka Topics. It requires more code but gives us flexibility in processing. However, it's not the main recommendation; using StreamingDataFrame is preferred.

• I needed to do a Delta Merge on the Delta Lake table, which the quixstreams library doesn't support directly. So, I used the low-level client API to achieve it.

• I like the quixstreams library because it offers a high-level API -StreamingDataFrame, for batch-like processing on streaming data. But if that's not an option, it also provides a low-level Kafka client API, allowing us to do almost anything.

Note: I have open-sourced my project & you can find all the source code here - Akashdesarda/data-pipeline-app-demo

Conclusion:

• Supercharging your streaming data pipeline in Python involves using modern tools and practices to boost efficiency and performance.

• Draw inspiration from batch processing to apply similar principles to streaming data without losing real-time capabilities.

• Use advanced libraries like Quix Streams for efficient data processing.

• Features like StreamingDataFrame these allow for high-level operations.

• Low-level client APIs are available for custom tasks.

• This approach simplifies development and provides flexibility and scalability.

• It enables effective tracking and analysis of customer journeys in real-time.

• Embracing these modern techniques ensures your data pipeline is robust, efficient, and meets the demands of today's data-driven environments.

Generally, the use of Data in any application can be categorised as

1. Analytical

2. ETL/ELT

Using SQL for analytical purposes is great, and you should keep using it. SQL truly excels in this area. There's no need for setup or worrying about the environment—just run the query and get your results.

But for ETL pipelines with complex transformations (which is common in the real world), you should switch from SQL immediately. Because using SQL often leads to very long, nested, and endless CTEs. This makes it really hard to read & nightmare debugging experience.

Challenges of SQL-heavy ETL Pipelines (Why it can be a mistake):

• Increased Complexity and Readability: ETL pipelines with many transformations often result in complex, deeply nested SQL queries with numerous CTEs, making the code hard to read, understand, and maintain. This complexity can be challenging for new team members or even for your future self to understand it.

• Debugging Nightmares: Tracking down errors in a massive SQL script with multiple levels of nesting can be a nightmare. The error messages might not always be as informative as those you'd get in your primary programming language's debugger.

• Limited Expressiveness for Complex Logic: While SQL is powerful for set-based operations, implementing intricate business logic or conditional transformations can become cumbersome and less intuitive compared to using the control flow structures available in languages like Python or Java.

• Version Control and Collaboration Challenges: While SQL scripts can be version-controlled, the changes and their impact might be harder to visualize and collaborate on compared to code within a larger application codebase managed by standard development tools and practices.

• Fragility and Maintainability: Complex SQL ETL can become fragile and hard to maintain over time, which might lead to constantly fixing problems as they arise.

• Testing Difficulties: Unit testing individual components and transformations within a large SQL script can be challenging. Testing often involves running the entire pipeline or significant portions of it, making it slower and less granular than testing functions or methods in your application code.

Below is a sample of a typical Query structure for an ETL pipeline. As you can see, even for fewer lines of code, it became complex to understand.

WITH
    SourceData AS (
        -- Step 1: Extract data from the raw source
        SELECT
            column1,
            column2,
            column3,
            -- ... more columns ...
            CAST(timestamp_column AS TIMESTAMP) AS processed_timestamp
        FROM
            raw_data_table
        WHERE
            date_column >= 'some_start_date' AND date_column < 'some_end_date'
    ),

    CleanedData AS (
        -- Step 2: Clean and standardize the data
        SELECT
            UPPER(TRIM(column1)) AS standardized_column1,
            CASE
                WHEN column2 = 'value_a' THEN 'Category A'
                WHEN column2 = 'value_b' THEN 'Category B'
                ELSE 'Other'
            END AS categorized_column2,
            CAST(column3 AS INTEGER) AS numeric_column3,
            processed_timestamp
        FROM
            SourceData
        WHERE
            column1 IS NOT NULL AND column3 IS NOT NULL
    ),

    TransformedData1 AS (
        -- Step 3: Apply the first set of transformations
        SELECT
            standardized_column1,
            categorized_column2,
            numeric_column3 * 10 AS transformed_numeric_column3,
            EXTRACT(YEAR FROM processed_timestamp) AS processing_year,
            EXTRACT(MONTH FROM processed_timestamp) AS processing_month
        FROM
            CleanedData
    ),

    LookupTable1 AS (
        -- Step 4: Join with a lookup table
        SELECT
            td1.*,
            lt1.lookup_value
        FROM
            TransformedData1 td1
        LEFT JOIN
            lookup_table_one lt1 ON td1.standardized_column1 = lt1.lookup_key
    ),

    FilteredData AS (
        -- Step 5: Apply specific filtering criteria
        SELECT
            *
        FROM
            LookupTable1
        WHERE
            processing_year = 2024 AND categorized_column2 IN ('Category A', 'Category B')
    ),

    AggregatedData AS (
        -- Step 6: Perform aggregations
        SELECT
            processing_year,
            processing_month,
            categorized_column2,
            COUNT(*) AS record_count,
            SUM(transformed_numeric_column3) AS total_transformed_value
        FROM
            FilteredData
        GROUP BY
            processing_year, processing_month, categorized_column2
    ),

    FinalStage AS (
        -- Step 7: Final transformations and calculations
        SELECT
            ad.*,
            (total_transformed_value / record_count) AS average_transformed_value,
            'Processed' AS status
        FROM
            AggregatedData ad
        -- Potentially another join with another lookup table here
        LEFT JOIN
            final_lookup_table flt ON ad.categorized_column2 = flt.category
    )

-- Step 8: Load or select the final processed data
SELECT
    *
FROM
    FinalStage
ORDER BY
    processing_year, processing_month, categorized_column2;

-- You might even have subqueries within these CTEs for more complex logic:
--
-- TransformedData2 AS (
--     SELECT
--         fd.*,
--         (SELECT MAX(another_column) FROM another_table WHERE fk_column = fd.some_id) AS max_value_from_other_table
--     FROM
--         FilteredData fd
-- ),
--
-- And so on...

The Solution

You should follow the below generic steps to solve the above-stated problem & improve your development experience.

• Step 1: Initial Starting Point

  • Use an ORM to pull decently filtered initial source data.

  • If you wish, you can even use Raw SQL with the programming language-specific database drivers.

  • If you are using Python, then I have already written a blog that explains how we should effectively work with databases in Python. It explains the solution which is the best of both world. You read here - How to effectively work with Databases in Python.

• Step 2: Design your codebase using a proper Library Structure

  • I'm confident that in the real world, you'll have many ETL pipelines running in a project to meet different business needs. I'm also certain that these pipelines will share a lot of common business logic.

  • All the code, business logic, utilities, connections, visualizations, and more should be developed and organized as a Library. This approach provides a wide range of flexible options and enables us to leverage the best practices of Software Engineering. You can apply principles like DRY, SRE, design patterns, OOP, and more. The possibilities are extensive.

  • Below is a sample library folder structure. This allows as t write relevant code to appropriate location.

  •   some_library/
      ├── config
      ├── utils/
      │   ├── module1
      │   └── module2
      ├── ops/
      │   ├── preprocessing/
      │   │   ├── module1
      │   │   └── module2
      │   └── processing/
      │       ├── module1
      │       └── module2
      └── pipeline/
          ├── pipeline1
          ├── pipeline2
          └── pipeline3
      cicd.yaml
      pyproject.toml
      poetry.lock
      test/
      ├── test_utils
      ├── test_ops
      └── test_pipeline
    

• Step 3: Use your library as a package in the pipeline

  • Now that the entire codebase is organized as a library, the next step is to publish it as a package.

  • Install and use this package in downstream pipelines just like any other open-source package.

  • This approach allows you to reuse the same functions for common functionalities.

• Step 4: Use the programming language ecosystem

  • Now that we're using programming languages like Python, Java, or JavaScript, we can take advantage of open-source tools and packages to enhance the pipeline even more.

  • This saves us from the pain of re-inventing the wheel.

Possible Counterarguments or Nuances to Consider:

• Performance: For very large datasets and simple, set-based transformations, well-optimized SQL can sometimes outperform code-based ETL due to the database's ability to leverage indexing and other performance optimizations. However, for complex logic, the overhead of repeated SQL calls and data transfer might negate this benefit.

• Database-Specific Features: Sometimes, specific database features might be most efficiently utilized directly through SQL. However, if the core logic becomes overly convoluted, the trade-off for readability and maintainability might still favor the above approach for the bulk of the pipeline.

• Out of Memory Issues: Large datasets won’t fit into memory & restrict us to run the pipeline. This can be solved using batch processing. To further improve performance, all the batches can be executed in parallel using a distributed system.

Conclusion:

The Situation

The project involved generating Monte Carlo Projections using Geometric Brownian Motion (GBM) to create thousands of potential future price paths for assets. This helps analysts estimate the probabilities of various outcomes.

First, I would like to provide a concise explanation of the components involved in this pipeline.

Pipeline Components

Step 1: Generate a random normal distribution and then perform a cumulative sum over it.

Step 2: Using the Covariance matrix and two sets of random normal distributions, compute the Einstein summation.

Step 3: Incorporating both deterministic trends (drift) and random fluctuations (volatility) using GBM.

Step 4: Perform mean aggregation & select based on filter criteria.

Issue

• The initial version used Numpy & Pandas library for all computations.

• We have to run the Projection for at least 100 years, but many times it would be even more than that. Let’s break down how the computation looks based on the above steps,

•     import numpy as np
  
      rng = np.random.default_rng()
  
      DAYS = 365
      YEARS = 1_00
      PATH = 20_000
  
      # Generating random cumulative sum normal distrubution
      # NOTE - We required two two normal distrubution set
      nor_dis_random1 = np.cumsum(
          rg.normal(0, np.sqrt(1 / days), (days * fixed_years, paths)),
          axis=0
      )
      nor_dis_random2 = np.cumsum(
          rg.normal(0, np.sqrt(1 / days), (days * fixed_years, paths)),
          axis=0
      )
  
      # Computing Einstein Summation
      ein_corr = np.einsum(
                      "ij,jkl->ikl",
                      covariance_matrix_data,
                      np.array([nor_dis_random1, nor_dis_random2]),
                      casting="same_kind",
                      optimize=True,
                  )
  

• As you can see above, we had to deal with extremely big arrays. Numpy has to allocate memory for random normal distribution generation. Business logic requires two such big arrays to compute Einstein Summation. For 100 years & above even after beefing it till 64 GB it was failing with Out of Memory Error.

• Following the Einstein Summation task, numerous transformations, including the application of GBM, were performed. This was initially done using Pandas , which made already memory hogging code even more bulky.

•     import pandas as pd
  
      # Creating two pandas dataframe from Einstein Summation aaray
      projecttion_a = pd.DataFrame(ein_corr[0].T).cumsum(axis=0)
      projecttion_b = pd.DataFrame(ein_corr[1].T).cumsum(axis=0)
  
      # Further transformation
      sigma_s = 0.26
      mu_s = 0.05
      s0 = 391555
  
      projecttion_a = S0 * np.exp((mu_S - 0.5 * sigma_S**2) + sigma_S * projecttion_a
      projecttion_b = S0 * np.exp((mu_S - 0.5 * sigma_S**2) + sigma_S * projecttion_b
  

• It became simply unsustainable to keep increasing memory.

The Tasks

Let me show you how I divided the problem into smaller tasks and solved them, aiming for a cohesive and efficient process.

Assessment

• The moment I stepped into this project, I realized that Pandas wasn't the right tool for the job. Both Numpy and Pandas store everything in memory from start to finish, and every transformation adds to memory use.

• While exploring possible solutions, I wondered about using PySpark because of its ability to handle distributed workloads. But then, I stumbled upon two significant issues:

  • It doesn't have first-class support for running NumPy functions across a cluster.

  • The environment is costly and bulky since Spark is cluster-based.

• Switching to Spark would have required a lot more work. That's where Polars comes in to save the day. Here are the key reasons why I chose it:

  • Like PySpark’s lazy evaluation, Polars also supports it with LazyFrame.

  • First-class for Numpy functions, even in LazyFrame.

  • Runs on a single machine and generally offers excellent performance since its based on Rust & internal use of parallel processing.

Redesigned Pipeline components

• Step 1: Replace all Pandas Dataframe code with Polars LazyFrame.

• Step 2: Since Polars fully supports Numpy functions, keep using specific Numpy functions that Polars doesn't have built-in.

• Step 3: Switch to Batch Processing for transformations. Ensure each batch creates a LazyFrame, so nothing is stored in memory until the final execution.

• Step 4: Keep adding transformation steps to the LazyFrame and execute them at the end. This approach lets Polars excel by making the most of its Lazy API.

With the lazy API, Polars doesn't run each query line-by-line but instead processes the full
query end-to-end. To get the most out of Polars it is important that you use the lazy API because:
 - The lazy API allows Polars to apply automatic query optimization with the query optimizer.
 - The lazy API allows you to work with larger than memory datasets using streaming.
 - The lazy API can catch schema errors before processing the data.

The Actions

Enough theory talking, now lets dive into real coding based actions that solved the all the above issues.

Numpy function + Polars Lazyframe

import polars as pl
import numpy as np

rng = np.random.default_rng()

DAYS = 365
YEARS = 1_00
PATH = 20_000

def unit_batch(fixed_years=10):
        return (
            pl.LazyFrame()
            .with_columns(
                # Here I am directly using NumPy's normal() & cumsum() in Polars Lazyframe since
                # it have first class support
                nor_dis_random=np.cumsum(
                    rg.normal(0, np.sqrt(1 / days), (days * fixed_years, paths)).astype(np.float32),
                    axis=0,
                    dtype=np.float32,
                ) # I am using Float32 as data structure instead of default Float64
            ) 
            # I need to create lazyframe of shape (x,y) from normal distribution data of shape (x,y)
            # Initially `nor_dis_random` is ArrayType column to which I am exploding to create y columns
            # The beauty of Polars Lazy API is we can keep adding steps to Query plan.
            .with_columns(pl.col("nor_dis_random").arr.to_struct().alias("array_struct"))
            .unnest("array_struct")
            .drop("nor_dis_random")
        )

 # Creating stack of all batches of lazyframe, that will eventually get concated
normal_dist_random_cum_sum_frames = [
    _unit_chunk(i)
    for i in alive_it(
        chunks, title="Normal Distribution Cumulative sum", force_tty=True, total=len(chunks)
    )
]
# This steps is very important as here I am vertically stacking all batches but still as LazyFrame
normal_dist = pl.concat(normal_dist_random_cum_sum_frames, how="vertical")

Batch Processing on Polars Lazyframe

BATCH_SIZE = 1_000

def calculate_no_of_batches(row_count: int, batch_size: int = 10) -> int:
    return row_count // batch_size + (row_count % batch_size > 0)


ein_corr_a, ein_corr_b = [], []

no_of_batches = calculate_no_of_batches(
    normal_dist1.select(pl.len()).collect().item(), batch_size=BATCH_SIZE
)

# Batch processing loop
for batch in alive_it(range(no_of_batches), title="Einstein summation", force_tty=True):
    # Polars Lazy API provides as with slice() to iterate over lazyframe without actually 
    # loading data in memory. Its similar to SQL Order & limit without actually loading data
    brownian_path_array1 = normal_dist1.slice(batch * BATCH_SIZE, BATCH_SIZE)
    brownian_path_array2 = normal_dist2.slice(batch * BATCH_SIZE, BATCH_SIZE)

    # Append as array to calculate Einstein summation convention
    einstein_summation_operands = np.array(
        [
            brownian_path_array1.collect().to_numpy().T,
            brownian_path_array2.collect().to_numpy().T,
        ],
        dtype=np.float32,
    )
    ein_corr = np.einsum(
        "ij,jkl->ikl",
        covariance_matrix_data,
        einstein_summation_operands,
        dtype=np.float32,
        casting="same_kind",
        optimize=True,
    )

    # Creating lazyframe from above & perform cumulative sum over entire column
    # Once again Polars Expressive Lazy API makes code more cleaner & readable. 
    ein_corr_a.append(pl.LazyFrame(ein_corr[0].T).with_columns(pl.all().cum_sum()))
    ein_corr_b.append(pl.LazyFrame(ein_corr[1].T).with_columns(pl.all().cum_sum()))

# Just like above, even here at end all smaller batches of lazyFrame is 
# concated as Unified Lazyframe. Still nothing is stored in memory.
ein_corr_frame1 = pl.concat(ein_corr_a, how="vertical")
ein_corr_frame2 = pl.concat(ein_corr_b, how="vertical")

Building data transformation Query graph

1. Computing GBM paths

# Get column names from standard_brownian for explicit transformations
col_names = ein_corr_frame1.collect_schema().names()

ein_corr_frame1 = (
    ein_corr_frame1
    # Step 1: Calculate drift term + volatility term for each column
    .with_columns(
        [
            (
                sigma_s * pl.col(col_name) + (mu_s - 0.5 * sigma_s**2) * pl.col("real_number_time")
            ).alias(f"{col_name}_gbm")
            for col_name in col_names
        ]
    )
    # Step 2: Apply exponential function and scale by s0
    .with_columns(
        [
            (pl.col(f"{col_name}_gbm").exp() * s0).alias(f"Path_{i + 1}")
            for i, col_name in enumerate(col_names)
        ]
    )
    # Step 3: Keep only the final sales path columns
    .select([f"Path_{i + 1}" for i in range(len(col_names))])
)

2. Compute Mean from GBM Paths & Use of Dynamic Scaling to fix Infinite number issue

# Get column names from standard_brownian for explicit transformations
col_names = ein_corr_frame1.collect_schema().names()

ein_corr_frame1 = (
    ein_corr_frame1.with_row_index("_idx", 1)
    # select row no present in year_end_location_time
    .filter(
        pl.col("_idx").is_in(
            time_step.select("year_end_location_time").drop_nulls().collect().to_series()
        )
    )
    .drop("_idx", strict=False)
    # Handle infinities by replacing them with None
    .with_columns(
        [
            pl.when(pl.col(c).is_infinite()).then(None).otherwise(pl.col(c)).alias(c)
            for c in col_names
        ]
    )  # Calculate max value per row for Dynamic Scaling ops
    .with_columns(max_value=pl.max_horizontal(pl.all()))
    # Performing dynamic scale mean by scaling down --> mean --> scale up
    .with_columns(
        Sales=pl.mean_horizontal(  # taking mean of entire row
            # dividing all cols by max value (scaled down)
            pl.all().exclude("max_value") / pl.col("max_value")
        )  # multiplying mean scaled down with max value (scale up)
        * pl.col("max_value")
    )  # We only need 'Sales' column
    .select("Sales")
)

3. Materializing the entire query plan into final output

# Till now we have performed many heavy maths based calculations, but everything is defred
# running collect() will materialize all the queries into final dataframe
ein_corr_frame1.collect()

Key Points:

• Polars expressive Lazy API is very powerful, clean & highly readable.

• From above you can easily gather that I am simple keep adding steps to Query plan without actually running them yet.

• This gives Polars’ engine many opportunities to optimize the query.

• Untill we don’t run LazyFrame.collect() all the queries are deferred & nothing is stored in memory

The Results

• The entire Monte Carlo Simulation running for 100 years were able to complete under 20-25 GB of memory, which was simply failing in using Pandas even after providing 64 GB.

• Added benefit was total time because Polars is based on Rust & internally uses parrallel processing to run queries.

• The Expressive API of Polars library is very powerful & intuitive, especially for Data engineers coming from SQL world.

• This experience underscores the importance of choosing the right tools and approaches in data engineering to achieve optimal performance and efficiency.

The age-old debate on the use of Raw SQL v/s ORM is still very much alive in today’s world. Let’s see some of the comparing points

A seasoned database expert might claim that an ORM or a programming language isn't necessary for working with a database. However, in practice, this approach has significant downsides, such as:

• For any complex problem, you might end up with multiple nested SELECT queries, which can be really tricky for others to debug and understand.

• SQL has a fixed set of keywords, so you have to work within those limits. This means missing out on all the amazing possibilities that a full-stack programming language like Python offers.

• You can't build a data pipeline using just SQL.

Therefore, the use of SQL should generally be limited to performing data analytics rather than developing data applications.

The World of ORMs

Python being a popular programming language provides many options for ORM. Let’s see them in action

Tech stack

• For this article, I am using Postgres database & Python 3.13

• Typically, an ORM transforms SQL data into Python data structures, which are then organized into either a dict, a tuple (a list of tuples), or even a namedtuple (a list of namedtuples).

• These basic data structures may not simplify handling real-world problems.

• I prefer using a Dataframe as the data structure to store the data. A Dataframe closely resembles a Database Table.

• You can choose between Pandas or Polars. Personally, I use Polars. I won't dive into why I prefer Polars over Pandas here, as that's a topic for another time. But trust me, switching to Polars can really make your life easier!

Benchmark

The environment that I use includes:

• Postgres 17

• Python 3.13

• Table of size 1,977,823

• This is the query

SELECT
    *
FROM
    FACTOR_INVESTING.TICKER_HISTORY
WHERE
    TICKER IN ('INFY', 'TCS')
    AND DATE BETWEEN CAST('2010-01-01' AS DATE) AND CAST('2024-01-01' AS DATE)
ORDER BY
    DATE DESC NULLS LAST

• The Benchmark query returns data of size 6910 rows

1. SQLAlchemy + Polars

from datetime import date
from about_time import about_time
from alive_progress import alive_it
from sqlalchemy.orm import declarative_base, mapped_column, Mapped, Session
from sqlalchemy import String, Date, DOUBLE_PRECISION, select, desc, create_engine
import polars as pl


engine = create_engine('postgresql+psycopg://akash:0330@localhost/playground')
Base = declarative_base()

class TickerHistory(Base):
    __tablename__ = "ticker_history"
    __table_args__ = {"schema": "factor_investing", "extend_existing": True}

    date: Mapped[date] = mapped_column(Date)
    ticker: Mapped[str] = mapped_column(String)
    key: Mapped[str] = mapped_column(String, primary_key=True, nullable=False)
    open: Mapped[float] = mapped_column(DOUBLE_PRECISION)
    high: Mapped[float] = mapped_column(DOUBLE_PRECISION)
    low: Mapped[float] = mapped_column(DOUBLE_PRECISION)
    close: Mapped[float] = mapped_column(DOUBLE_PRECISION)

query = (
    select(TickerHistory)
    .where(TickerHistory.ticker.in_(["INFY", "TCS"]))
    .where(TickerHistory.date.between(date(2010, 1, 1), date(2024, 1, 1)))
    .order_by(desc(TickerHistory.date))
)

with about_time() as t:
    with Session(engine) as session:
        # Running the same query 100 times
        for i in alive_it(range(100)):
            # Directly reading query in polars dataframe
            df = pl.read_database(query, session)

print(f"Total time taken: {t.duration_human}")

The result

So that’s 44.41 second for 100 queries & 2.25 queries per second throughout

2. Pony ORM, Peewee ORM, SQLModel ORM

Both of these popular ORMs currently do not support Psycopg3 and still require Psycopg2. The library authors recommend using Psycopg3 going forward, as mentioned here. Because of this, I decided to skip using both of these ORMs.

SQLModel ORM which is itself based on SQLAlchemy, so we won't any significant difference in results.

But this is not the end of the tunnel. There are other potentially good options too. Let’s check them out

You may have already got a feeling that I prefer ORM based solution compared to Raw SQL. But theoretically, Raw SQL query execution should have an edge over ORM execution.

To combine the power of both the world I use Sqlglot Library. This allows you to build the query programmatically as well as iteratively. Eliminates issues like typos/spelling errors, SQL Injection, etc. It supports 24 dialects.

3. Psycopg3 + Sqlglot + Polars

from datetime import date
from about_time import about_time
from alive_progress import alive_it
from psycopg import connect
import polars as pl
from sqlglot import select, condition, Dialects

query = (
    select("*")
    .from_("factor_investing.ticker_history")
    .where(condition("ticker").isin("INFY", "TCS"))
    .where(condition("date").between(date(2010, 1, 1), date(2024, 1, 1)))
    .order_by("date DESC")
    .sql(Dialects.POSTGRES)
)

conn = connect(
    host="localhost", port=5432, dbname="playground", user="akash", password="0330"
)

with about_time() as t:
    print("Starting benchmark...")
    with conn.cursor() as cursor:
        # Running the same query 100 times
        for i in alive_it(range(100)):
            # Directly reading query in polars dataframe
            df = pl.read_database(query, cursor)

print(f"Total time taken: {t.duration_human}")

So that’s 20.04 second for 100 queries & 5 queries per second throughout. This improves the result compared to SQLAlchemy.

4. ADBC + Sqlglot + Polars

Polars dataframe are backed by Arrow Table & uses PyArrow. Use of ADBC (Arrow Database Connectivity) can benefit from zero copy concept, since Polars doesn't need to convert table data returned by ADBC drivers.

from adbc_driver_postgresql.dbapi import connect
from datetime import date
from about_time import about_time
from alive_progress import alive_it
import polars as pl
from sqlglot import select, condition, Dialects

query = (
    select("*")
    .from_("factor_investing.ticker_history")
    .where(condition("ticker").isin("INFY", "TCS"))
    .where(condition("date").between(date(2010, 1, 1), date(2024, 1, 1)))
    .order_by("date DESC")
    .sql(Dialects.POSTGRES)
)

uri = "postgresql://akash:0330@localhost/playground"

with about_time() as t:
    with connect(uri) as conn:
       # Running the same query 100 times
        for i in alive_it(range(100)):
            # Directly reading query in polars dataframe
            df = pl.read_database(query, conn)

print(f"Total time taken: {t.duration_human}")

So that’s 17.94 second for 100 queries & 5.68 queries per second throughout. This improves the result compared to both SQLAlchemy & Psycopg

5. ConnectorX + Sqlglot + Polars

ConnectorX is yet another SQL driver/library that is making some rounds. It is using Rust. Let’s see in action. BTW it uses a slightly different approach. There is no connection or cursor concept here.

 from datetime import date
from about_time import about_time
from alive_progress import alive_it
import polars as pl
from sqlglot import select, condition, Dialects

query = (
    select("*")
    .from_("factor_investing.ticker_history")
    .where(condition("ticker").isin("INFY", "TCS"))
    .where(condition("date").between(date(2010, 1, 1), date(2024, 1, 1)))
    .order_by("date DESC")
    .sql(Dialects.POSTGRES)
)

uri = "postgresql://akash:0330@localhost/playground"
with about_time() as t:
    # Running the same query 100 times
        for i in alive_it(range(100)):
            # Directly reading query in polars dataframe
            df = pl.read_database_uri(query,uri, engine="connectorx")

print(f"Total time taken: {t.duration_human}")

So that’s 1:22.4 second for 100 queries & 12.1 queries per second throughout. This is the worst of the lot.

Conclusion

• The thought that Raw SQL can give you better results holds true. But at the same time directly using it in Python will be never a good idea.

• Sqlglot nicely fits this scenario. Allows us to use Raw SQL query & combine with Dataframe.

• In terms of benchmarks, the combination of ADBC + Sqlglot + Polars is the winner because of the tight Arrow Integration. So this should be your first pick.

• But at the same time sticking with Psycopg is also not a bad idea.

• Design and develop business logic.

• A notebook that performs all the business logic.

• Running that notebook using Databricks Workflow.

This is easier said than done. Logistics for running this notebook(s) is one of the biggest headaches. I am sure you won’t be just running one or two notebooks. Setting the environment correctly for every one of them can be cumbersome. It becomes even tougher if you are using proprietary/private libraries.

Following is the system that I came up with which is the most practical solution:

1. Step 1: Universal environment setup notebook.

   • Create a universal environment setup notebook in your repo.

   • This notebook can be placed in the root of your repo so that all other notebooks can easily access it

2. Step 2: Secrets, Environment variables, Constants, etc

   • Secrets, variables, and constants change depending on the environment in which the notebook is running, so they must be set accordingly.

   • We use dbutils.secret to fetch all the secrets. Let’s see in action

       # Getting all secrets
       # Note - It is important to make sure that the scope is appropriately mapped to
       # the secret store assigned to environment-specific workspace
       PAT = dbutils.secrets.get("<your_secret_scope>", "<secret_1_name>")
       SECRET_1 = dbutils.secrets.get("<your_secret_scope>", "<secret_1_name>")
       SECRET_2 = dbutils.secrets.get("<your_secret_scope>", "<secret_2_name>")
       SECRET_3 = dbutils.secrets.get("<your_secret_scope>", "<secret_3_name>")
     

       # Environment specific constants
       databricks_host = spark.conf.get("spark.databricks.workspaceUrl")
     
       if databricks_host == "<your dev workapce host url>":
           environment = "dev"
           CONSTANT_1 = "dev constant 1"
           CONSTANT_2 = "dev constant 2"
           volume_path = "/Volumes/dev/path"
           catalog_name = "dev_catalog"
       elif databricks_host == "<your uat workapce host url>":
           environment = "uat"
           CONSTANT_1 = "uat constant 1"
           CONSTANT_2 = "uat constant 2"
           volume_path = "/Volumes/uat/path"
           catalog_name = "uat_catalog"
       else databricks_host == "<your prod workapce host url>":
           environment = "prd"
           CONSTANT_1 = "prod constant 1"
           CONSTANT_2 = "prod constant 2"
           volume_path = "/Volumes/prod/path"
           catalog_name = "prod_catalog"
       else:
           raise NameError("Incorrect databricks workspace")
     

3. Step 3: Passing all variables to the downstream notebookThere are a couple of ways to pass these variables to the downstream notebook to access environment-specific values. Here are two methods I use:

   • DBFS JSON file: Save all the values in a JSON file at the dbfs/tmp/ location. This way, the file is deleted once the notebook job is done. I used this method for a while, but Databricks recently changed the permission rules, and now only admins can access it.

   • Temp View: Using a TEMP VIEW is even better than a file. It doesn't require admin permission. So, create a TEMP VIEW that includes all these variables.

•   # Creating dict of all variables
  
    env_vars = {
        "secret_1": SECRET_1,
        "secret_2": SECRET_2,
        "secret_3": SECRET_3,
        "constant_1": CONSTANT_1,
        "constant_2": CONSTANT_2,
        "catalog_name": catalog_name
    }
  
    # Write environment variables to a TEMP VIEW
    spark.createDataFrame([env_vars]).createOrReplaceTempView("env_vars")
  

4. Step 4: Installing libraries, especially private

   • You should definitely package your codebase as a Python Library. Then install & use it just like any other open source library. This way you won’t have to worry about Path or Relative/Absolute Import Error, etc.

   • You need to separate package for dev/qa (which can be experimental) & uat/prod (must be stable)

   • Follow PEP 440 guidelines to version your code.

     • use X.Y.devN version for the package published from the develop branch.

     • use X.Y.N version for the package published from the master branch.

   • Use %pip to install environment-specific private library.

       if environment == "dev" or environment == "qa":
           # --pre flag will install package having 'dev' label
           # NOTE - use your repository usrl appropriately
            %pip install --extra-index-url https://databricks-user:{PAT}@pkgs.dev.azure.com/dp-us-bus/_packaging/dp-pip/pypi/simple/ --pre --upgrade "<your package name>"
     
       if environment == "uat" or environment == "prod":
           %pip install --extra-index-url https://databricks-user:{PAT}@pkgs.dev.azure.com/dp-us-bus/_packaging/dp-pip/pypi/simple/ --upgrade "<your package name>"
     
       # Note - Databricks recommends to restart python to make sure we'll 
       # be using libraries that were just installed
       dbutils.restartPython()
     

5. Step 5: Running the setup notebook in the downstream job notebook

   • Use databricks’ magic %run command.

   • Read the TEMP VIEW & update values in the os environment

       # Running the notebook. Make sure to use correct relative path
       %run ../../databricks_environment_setup
     
       # Read environment variables from the TEMP VIEW & set the environment variables
       # for use in this notebook
       import os
     
       os.environ.update(spark.table("env_vars").first().asDict())
     

Here are the steps

1. Step 1: Use FastAPI-Auth0 python library.

2. Step 2: Create Auth0 FastAPI Security dependency

from typing import Any, Iterator

from fastapi import Depends, FastAPI, Security
from fastapi_auth0.auth import Auth0
from pydantic import RootModel


class RootUser(RootModel):
    root: dict[str, Any]

    def __iter__(self) -> Iterator[str]:
        return iter(self.root)

    def __getitem__(self, item) -> Any:
        return self.root[item]

auth = Auth0(
    domain="<Your auth0 domain>",
    api_audience="<Your auth0 api audience>",
    auth0user_model=RootUser
) # Note - You can use customized pydantic user model too

3. Step 3: Secured the Endpoint using FastAPI dependency

app = FastAPI()

# NOTE - I am using `auth.implicit_scheme` to implement Auth0 SSO
@app.get("/protected", dependencies=[Depends(auth.implicit_scheme)])
# NOTE -`auth.get_user` performs the login & returns user identity 
def get_secure(user: RootUser = Security(auth.get_user)):
    return {"message": f"{user}"}

4. Step 4: (Optional) If you want the client id to be auto-filled then modify it accordingly

from typing import Any, Iterator

from fastapi import Depends, FastAPI, Security
from fastapi_auth0.auth import Auth0
from pydantic import RootModel


class RootUser(RootModel):
    root: dict[str, Any]

    def __iter__(self) -> Iterator[str]:
        return iter(self.root)

    def __getitem__(self, item) -> Any:
        return self.root[item]


auth = Auth0(
    domain="<Your auth0 domain>",
    api_audience="<Your auth0 api audience>",
    auth0user_model=RootUser
) # Note - You can use a customized pydantic user model too
app = FastAPI(swagger_ui_init_oauth={"clientId": "<your client id>"})

@app.get("/protected", dependencies=[Depends(auth.implicit_scheme)])
def get_secure(user: RootUser = Security(auth.get_user)):
    return {"message": f"{user}"}

They are incredibly useful when you need to:

• Share logic effortlessly (use the same code logic over and over again).

• Seamlessly share database connections.

• Enforce security, authentication, role requirements, and more with ease.

• And so many other amazing things...

All this, while keeping code repetition to a minimum!

I started using Dependency Injection just for authentication. So I used it like this

from .models import (
    Auth0APIUser,
    RootAuth0User,
)

# Security flow to authenticate user using Password flow
header_username = APIKeyHeader(
    name="username", scheme_name="username", auto_error=False
)
header_password = APIKeyHeader(
    name="password", scheme_name="password", auto_error=False
)

# Security flow to authenticate user using Auth0
auth0_authentication = Auth0(
    domain=env_config.auth0_domain,
    api_audience=env_config.auth0_api_audience,
    auth0user_model=RootAuth0User,  # type: ignore
    auto_error=False,
)

def user_authentication_security(
    user: Annotated[RootAuth0User, Security(auth0_authentication.get_user)],
) -> Auth0APIUser:
    """FASTApi Security injection to perform user auth using auth0"""
    return auth0_api_response_to_user(user.model_dump())

Next is its application in the API Endpoint function

@app.get("/schema/{schema}/tables/{table}")
async def _get_table_data(
    schema: str,
    table: str,
    user: Annotated[Auth0APIUser, Depends(user_authentication_security)]
) -> dict:
    # checking if users has access to schema
    schema_access = await check_schema_access(schema, user)
    if schema_access is False:
        logger.warning(
            f"{user_info.email_id} is not authorized to access {schema}",
            extra=logger_properties,
        )
        raise HTTPException(
            status.HTTP_403_FORBIDDEN,
            detail=f"{user_info.email_id} is not authorized to access {schema}",
        )

    # checking if users has access to table
    table_access = await check_table_access(schema, table, user)
    if table_access is False:
        logger.warning(
            f"{user_info.email_id} is not authorized to access {table} under {schema}",
            extra=logger_properties,
        )
        raise HTTPException(
            status.HTTP_403_FORBIDDEN,
            detail=f"{user_info.email_id} is not authorized to access {table} under {schema}",
        )

    # checking if table is present or not
    # checking if requested dataset is present
    current_dataset = await datasetDetailsV2.find_one(
        datasetDetailsV2.dataset == dataset
    )
    if not current_dataset:
        raise HTTPException(
            status.HTTP_404_NOT_FOUND,
            detail=f"dataset: {dataset} not found",
        )
    # checking if requested table is present in dataset
    complete_table_name = f"{current_dataset.destination_metadata.unity_catalog}.{current_dataset.dataset}.{table}"
    if complete_table_name not in current_dataset.delta_share_compatible_table:
        raise HTTPException(
            status.HTTP_404_NOT_FOUND,
            detail=f"table: {table} not found in dataset: {dataset}",
        )

    # From here onwards comes the core logic of above endpoint logic
    ...

The code above is practically shouting, "I need some help! 😫". This are several issues:

• A lot of other code is being used in the function apart from the core logic. This violates the Single Responsibility Principle (SRP) of SOLID Principal

• If I need to check if a Table is present or if a User has the right access in another endpoint function, I will end up having to copy and paste the same logic there too.

• What if I need to change the authentication and authorization logic in the future? It would make testing a lot more complicated too.

• Also, long functions just look too ugly 🤮

We can improve this by removing things like authentication, authorization, verification, etc. and all the non-core logic from the API endpoint function into a separate dependency function. We can even use the nested dependency function to comply with Don’t Repeat Yourself (DRY).

from .models import (
    Auth0APIUser,
    RootAuth0User,
)

# Security flow to authenticate user using Password flow
header_username = APIKeyHeader(
    name="username", scheme_name="username", auto_error=False
)
header_password = APIKeyHeader(
    name="password", scheme_name="password", auto_error=False
)

# Security flow to authenticate user using Auth0
auth0_authentication = Auth0(
    domain=env_config.auth0_domain,
    api_audience=env_config.auth0_api_audience,
    auth0user_model=RootAuth0User,  # type: ignore
    auto_error=False,
)

def user_authentication_security(
    user: Annotated[RootAuth0User, Security(auth0_authentication.get_user)],
) -> Auth0APIUser:
    """FASTApi Security injection to perform user auth using auth0"""
    return auth0_api_response_to_user(user.model_dump())

async def dataset_table_availability_verification(dataset: str, table: str):
    """FASTApi dependency to check if requested dataset and table is present"""
    # checking if requested dataset is present
    current_dataset = await datasetDetailsV2.find_one(
        datasetDetailsV2.dataset == dataset
    )
    if not current_dataset:
        raise HTTPException(
            status.HTTP_404_NOT_FOUND,
            detail=f"dataset: {dataset} not found",
        )
    # checking if requested table is present in dataset
    complete_table_name = f"{current_dataset.destination_metadata.unity_catalog}.{current_dataset.dataset}.{table}"
    if complete_table_name not in current_dataset.delta_share_compatible_table:
        raise HTTPException(
            status.HTTP_404_NOT_FOUND,
            detail=f"table: {table} not found in dataset: {dataset}",
        )
    return current_dataset

# 1st Level dependancy
async def role_verification(
    user: Annotated[Auth0APIUser, Depends(get_user)], claim: str
) -> userInfoV2:
    # getting current user by matching auth0 api user id
    current_user = await userInfoV2.find_one(userInfoV2.user_id == user.user_id)
    if not current_user:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="user not found"
        )
    # checking if requested role is present
    user_rbac = RBACPolicyClient(user_policy=current_user.model_dump())
    try:
        user_rbac.enforce_claim(policy_level=PolicyLevel.role, claim=claim)
    except RBACPolicyViolation as e:
        logger.warning(f"{user.primary_email} not approved for {claim}")
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail=f"{user.primary_email} don't have role: {claim}",
        ) from e

# 1st Level dependancy
async def scope_verification(
    user: Annotated[Auth0APIUser, Depends(get_user)], claim: str
) -> userInfoV2:
    # getting current user by matching auth0 api user id
    current_user = await userInfoV2.find_one(userInfoV2.user_id == user.user_id)
    if not current_user:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="user not found"
        )
    user_rbac = RBACPolicyClient(user_policy=current_user.model_dump())
    try:
        user_rbac.enforce_claim(policy_level=PolicyLevel.scope, claim=claim)
    except RBACPolicyViolation as e:
        logger.warning(f"{user.primary_email} not approved for {claim}")
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail=f"{user.primary_email} don't have scope: {claim}",
        ) from e
    return current_user

# 2nd Level dependency
async def write_table_scope_dependency(
    user: Annotated[Auth0APIUser, Depends(get_user)],
    dataset: str,
    table: str,
) -> userInfoV2:
    scope_condition = f"{dataset}:{table}::write"
    return await scope_verification(user=user, claim=scope_condition)
# 2nd Level dependency
async def read_table_scope_dependency(
    user: Annotated[Auth0APIUser, Depends(get_user)],
    dataset: str,
    table: str,
) -> userInfoV2:
    scope_condition = f"{dataset}:{table}::read"  # read scope will present by default
    return await scope_verification(user=user, claim=scope_condition)
# 2nd Level dependency
async def engineering_admin_role_dependency(
    user: Annotated[Auth0APIUser, Depends(get_user)],
) -> userInfoV2:
    return await role_verification(user=user, claim="management:engineering::admin")

Note that I've reused the 1st level dependency in several 2nd level dependencies, following the DRY principle. Up next, we'll see how it's applied in an API Endpoint function.

# EG 1. Here I am using authorization as dependancy rather in the same function
@app.get("/user", dependencies=[Depends(engineering_admin_role_dependency)])
async def _list_users() -> list[userInfoV2]:
    """Get the list of all users along with their access"""
    return await userInfoV2.find_all().to_list()

# EG 2. Similar to above, I used authorization as dependancy
@app.patch(
    "/dataset",
    status_code=status.HTTP_201_CREATED,
    dependencies=[Depends(engineering_admin_role_dependency)],
)
async def _update_dataset(details: DatasetInput):
    """Update existing dataset into the system"""

    # checking if dataset is present
    current_dataset = await datasetDetailsV2.find_one(
        datasetDetailsV2.dataset == details.dataset
    )
    if not current_dataset:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="dataset not present",
        )
    # Core logic will come here

# EG 3. Using multiple dependancy to make code more clearner
@router_v2.get(
    "/dataset/{dataset}/table/{table}",
    dependencies=[Depends(read_table_scope_dependency)],
    response_class=UJSONResponse,
)
async def _get_table_data(
    dataset: str,
    table: str,
    current_dataset: Annotated[
        datasetDetailsV2, Depends(dataset_table_availability_verification)
    ],
) -> list[dict]:
    """get delta table from desired dataset & table"""
    # Core logic to read table data will come here

The key takeaways we can gather from the above are:

• Keep things tidy by separating dependencies like auth and verification from the endpoint function. This makes the code look cleaner.

• Use multilevel dependency functions to reuse code that repeats.

• Testing becomes easier since we can create dedicated tests focusing on the endpoint function and other tests focusing on the dependency function.

• Other endpoints with the same needs can easily reuse the same dependency injections.

All of this helped cut down the lines of code, streamline the logic, and generally make the code easier to read.

But there's still room for improvement. Some logic doesn't really belong to the core logic. By building on the Multilevel Dependency Injection, we can further optimize it as follows:

# 1st Level dependency
async def user_presence_dependency(
    request: Request,
    user: Annotated[Auth0APIUser, Depends(get_user)],
) -> UserInfoV2:
    current_user_df = pl.read_database(
        current_user_filter(user.user_id), request.app.state.cursor
    )
    if current_user_df.is_empty():
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"{user.primary_email} not found",
        )
    return UserInfoV2.from_polars(current_user_df)

async def dataset_presence_dependency(
    request: Request, dataset: str,
) -> DatasetInfoV2:
    current_dataset_df = pl.read_database(
        current_dataset_filter(dataset), request.app.state.cursor
    )
    if current_dataset_df.is_empty():
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"{dataset} not found",
        )
    return DatasetInfoV2.from_polars(current_dataset_df)

# 2nd Level dependency
async def table_presence_dependency(
    dataset: Annotated[DatasetInfoV2, Depends(dataset_presence_dependency)],
    table: str,
) -> str:
    if dataset.table is None or table not in dataset.table:
        raise HTTPException(
            status.HTTP_404_NOT_FOUND,
            detail=f"table: {table} not found in dataset: {dataset.dataset}",
        )
    return table

# 3rd Level dependency
async def read_table_scope_dependency(
    user: Annotated[UserInfoV2, Depends(user_presence_dependency)],
    dataset: str,
    table: str,
) -> UserInfoV2:
    scope_condition = f"{dataset}:{table}::read"  # read scope will present by default
    return await scope_verification(user=user, claim=scope_condition)


async def engineering_admin_role_dependency(
    user: Annotated[UserInfoV2, Depends(user_presence_dependency)],
) -> UserInfoV2:
    return await role_verification(user=user, claim=FixedString.ENGINEERING_ADMIN.value)

Next is its application in the API Endpoint function

@router_v2.patch(
    "/dataset/{dataset}",
    status_code=status.HTTP_201_CREATED,
    dependencies=[Depends(engineering_admin_role_dependency)],
)
async def _update_dataset(
    dataset: Annotated[DatasetInfoV2, Depends(dataset_presence_dependency)],
) -> dict[str, str]:
    """Update the existing dataset details"""
    # Core logic will come here

async def _read_data(
    dataset: Annotated[DatasetInfoV2, Depends(dataset_presence_dependency)],
    table: Annotated[str, Depends(table_presence_dependency)],
    filter_query: Annotated[FilterQueryParams, Query()],
    request: Request,
) -> ORJSONResponse:
    """Get data from desired dataset & table"""
    # Core logic will come here

The key takeaways we can gather from the above are:

• The code now looks even cleaner. The API endpoint function now includes only the code related to itself and nothing else.

• You might think that adding more dependency functions means adding more lines of code, but these are reusable dependencies. When you have many endpoints with the same requirements, the benefits really add up!

In this article, I have used 2-3 endpoints for example. But my actual project has upwards of 30 endpoints. Following were the impact.

• They are grouped by categories. Ideally, all the endpoints in a category will have common requirements like authentication methods. I simply used the security dependency injection at the API Router definition & all the endpoint functions referring to it don’t need to use it again.

• Even for the endpoints which have additional requirements, I can just use the respective dependency.

• That's why overall I was able to reduce the code in my project by more than 30%

• Apart from code optimization, if you closely look at the endpoints having either path or query parameters, using dependencies becomes a lot more readable. Now anyone knows what the requirements are for that specific parameter.

Introduction

Recap of Part 1

In Part 1, we discussed why having a strong enterprise data strategy is important and the architectural choices we made to fit our needs.

A big part of an enterprise data strategy is data sharing. In this article, we'll dive into how we built our data sharing application, the tools and technology we used, and the valuable lessons we picked up along the way.

The need

Even before we jump in, I want to discuss what were the most important requirements that we were aiming to solve

• We chose the Databricks Lakehouse platform for data storage. It's based on the Spark ecosystem, using Spark SQL or PySpark for data interaction. However, most users and tools don't primarily use these methods.

• Most of our users are financial analysts who have long used Ms Excel for their workflows. We aimed to create a solution that securely provides them with authorized data access.

System Backend Architecture

To select the backend architecture we took inspiration from how two service communicate with each other - APIs. We explored multiple API Driven Architecture like REST API, GraphQL, gRPC, but ultimately decided to go ahead with REST API due to its widespread compatibility with tools (including Excel 😍)

We used FastAPI as the REST API framework.

Project Structure

├── data_sharing/
│   ├── dependency/
│   │   ├── auth.py
│   │   └── verification.py
│   ├── routers/
│   │   ├── admin.py
│   │   ├── users.py
│   │   └── table.py
│   ├── models.py
│   └── config.py
└── main.py

• We designed the structure of the project keeping modularity as a priority. Not everything is cramped into main.py. Follow this doc for more info.

• All the routers have their modules. Similarly, dependencies too have their modules.

• We also have a dedicated models module to declare all sorts Pydantic, Enum, Dataclass, etc models. Similarly a module for the environment config to make the codebase truly environment-agnostic without changing a single line of code.

API Endpoints

These are some of the important endpoints

1. GET /api/v1/users : To get a list of all user details

2. GET /api/v1/users/{user_id}: To get details of specific user_id.

3. PUT /api/v1/users/{user_id}/rbac: Override specific user_id RBAC with provided values

4. PATCH /api/v1/users/{user_id}/rbac?policy-level: Change the specific RBAC policy level of the given user_id.

5. DELETE /api/v1/users/{user_id}: Delete the given user_id from the system.

6. GET /api/v1/dataset: To get a list of all dataset details.

7. GET /api/v1/dataset/{dataset}: To get the details of the given dataset

8. GET /api/v1/dataset/{dataset}/table: To get the list of all tables in the given dataset

9. GET /api/v1/dataset/{dataset}/table/{table}: To get the table data residing in the given dataset

10. POST /api/v1/dataset/{dataset}/table/{table}: To write data as payload to the given table of the given datset

11. POST /api/v1/dataset/{dataset}/table/{table}/file: To write data as file upload using Form data protocol to the given table of the given datset

Authentication

The most important part of any authentication logic is getting the identity of the user. We are using Auth0 & SSO provided by it. Let’s see in action how to implement it in FastAPI.

• The biggest challenge that we had was to integrate SSO login with Swagger Docs. We started with the official docs provided by Auth0. This is more on the manual side for U2M OAuth flow.

• As I said earlier our goal was to use SSO login instead. So we instead use fastapi-auth0 python library which enabled us to use SSO login write into Swagger docs.

• To abstract all the complexity from end user we combined Security with Dependency Inject of FastAPI.

  💡

  FastAPI doesen’t really restricts us with Dependency Injection. We can go craazy creative using it. Its one of the best feature that FastAPI provides us.

# For Auth we need couple of models. This is written in data_sharing/modes.py

from pydantic import (
    BaseModel,
    EmailStr,
    Field,
    RootModel
)

class RootAuth0User(RootModel):
    root: dict[str, Any]

    def __iter__(self):
        return iter(self.root)

    def __getitem__(self, item):
        return self.root[item]

class Auth0APIUser(BaseModel):
    user_id: str
    primary_email: EmailStr
    sub: str
    aud: str
    iat: datetime
    exp: datetime
    azp: str

# This SSO Auth code is written in data_sharing/dependency/auth.py


from fastapi import Depends, HTTPException, Security, status
from fastapi.security import APIKeyHeader, HTTPBasic
from fastapi_auth0.auth import Auth0, JwksDict

from data_sharing.models import (
    Auth0APIUser,
    RootAuth0User,
)

# Security flow to authenticate user using Auth0
auth0_authentication = Auth0(
    domain="<your auth0 domain>",
    api_audience="<your auth0 api audience>",
    auth0user_model=RootAuth0User,  # type: ignore
    auto_error=False,
)

async def _get_user_using_auth0_sso_flow(
    user: Annotated[RootAuth0User, Security(auth0_authentication.implicit_scheme)],
) -> RootAuth0User:
    """FASTApi Security injection to perform user auth using auth0 sso flow."""
    return user


async def _get_user_using_auth0_implicit_flow(
    user: Annotated[RootAuth0User, Security(auth0_authentication.get_user)],
) -> RootAuth0User:
    """FastAPI Security injection to perform user auth using auth0 implicit (Bearer token) flow."""
    return user

# Once successfull SSO login & code exchange with Auth0, we have to convert it to actucal user
async def get_user(
    sso_user: Annotated[RootAuth0User, Depends(_get_user_using_auth0_sso_flow)],
    implicit_user: Annotated[
        RootAuth0User, Depends(_get_user_using_auth0_implicit_flow)
    ],
) -> Auth0APIUser:
    """FastAPI Security injection to get the user based on the various auth flow"""
    # NOTE - keep adding all the auth mechanism here to get the user
    if sso_user:
        # If you don't get response that follows AuthOAPIUser model then you
        # may neeed to writea small function to convert it, like below
        # return auth0_api_response_to_user(sso_user.model_dump())
        return Auth0APIUser(**sso_user.model_dump())
    elif implicit_user:
        # return auth0_api_response_to_user(implicit_user.model_dump())
        return Auth0APIUser(**sso_user.model_dump())
    else:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="No authorization token provided",
        )

Authorization

• We cannot treat authorization as some afterthought. It’s equally important. As stated earlier we use RBAC Strategy for authorization.

• It’s implementation is again powered by FastAPI Dependency. (See that’s why I say I love it ❤️)

# All this authorization & verification code is written in data_sharing/dependency/verification.py

import polars as pl

async def role_verification(user: UserInfo, claim: str) -> UserInfo:
    # checking if requested role is present
    user_rbac = RBACPolicyClient(user_policy=user.model_dump())
    try:
        user_rbac.enforce_claim(policy_level=PolicyLevel.role, claim=claim)
    except RBACPolicyViolation as e:
        logger.warning(
            f"{user.primary_email} not approved for {claim}"
        )
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail=f"{decrypt_data(ENCRYPTION_KEY,user.primary_email)} don't have role: {claim}",
        ) from e
    return user


async def scope_verification(user: UserInfo, claim: str) -> UserInfo:
    user_rbac = RBACPolicyClient(user_policy=user.model_dump())
    try:
        user_rbac.enforce_claim(policy_level=PolicyLevel.scope, claim=claim)
    except RBACPolicyViolation as e:
        logger.warning(
            f"{user.primary_email} not approved for {claim}"
        )
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail=f"{user.primary_email} don't have scope: {claim}",
        ) from e
    return user


async def permission_verification(user: UserInfo, permission: str) -> UserInfo:
    # checking if requested permission is present
    if user.permission is not None and permission not in user.permission:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail=f"user don't have permission: {permission}",
        )
    return user


async def user_presence_dependency(
    request: Request, user: Annotated[Auth0APIUser, Depends(get_user)],
) -> UserInfo:
    current_user_df = pl.read_database(
        current_user_filter(user.user_id), request.app.state.cursor
    )
    if current_user_df.is_empty():
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"{user.primary_email} not found",
        )
    return UserInfo.from_polars(current_user_df)


async def dataset_presence_dependency(request: Request, dataset: str
) -> DatasetInfo:
    current_dataset_df = pl.read_database(
        current_dataset_filter(dataset), request.app.state.cursor
    )
    if current_dataset_df.is_empty():
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"{dataset} not found",
        )
    return DatasetInfoV2.from_polars(current_dataset_df)


async def table_presence_dependency(
    dataset: Annotated[DatasetInfo, Depends(dataset_presence_dependency)], table:
) -> str:
    if dataset.table is None or table not in dataset.table:
        raise HTTPException(
            status.HTTP_404_NOT_FOUND,
            detail=f"table: {table} not found in dataset: {dataset.dataset}",
        )
    return table


async def write_table_scope_dependency(
    user: Annotated[UserInfo, Depends(user_presence_dependency)],
    dataset: str,
    table: str,
) -> UserInfo:
    scope_condition = f"{dataset}:{table}::write"
    return await scope_verification(user=user, claim=scope_condition)


async def read_table_scope_dependency(
    user: Annotated[UserInfo, Depends(user_presence_dependency)],
    dataset: str,
    table: str,
) -> UserInfo:
    scope_condition = f"{dataset}:{table}::read"  # read scope will present by default
    return await scope_verification(user=user, claim=scope_condition)


async def engineering_admin_role_dependency(
    user: Annotated[UserInfo, Depends(user_presence_dependency)],
) -> UserInfo:
    return await role_verification(user=user, claim=FixedString.ENGINEERING_ADMIN.value)


async def dataset_owner_role_dependency(
    user: Annotated[UserInfo, Depends(user_presence_dependency)], dataset: str
) -> UserInfo:
    return await role_verification(user=user, claim=f"{dataset}::dataset_owner")

Routers in action

Lets see few example how we can combine all the above to create API Routers.

# All user router code is written in data_sharing/router/users.py

# SECTION - User router
router = APIRouter(prefix="/api/v1", tags=[APITags.user])

@router.get("/user/{user_id}")
async def _get_user(
    user: Annotated[UserInfo, Depends(user_presence_dependency)],
    user_id: str
    request: Request,
) -> UserInfoV2:
    """Get the details of specified `user` along with their access"""
    # Your respective logic to get user details will come here

@router.put(
    "/user/{user_id}/rbac",
    status_code=status.HTTP_201_CREATED,
    dependencies=[Depends(engineering_admin_role_dependency)],
)
async def _replace_user_rbac_v2(
    user_id: str, user_rbac: UserRBACInfo
) -> dict[str, str]:
    """Replace user's complete RBAC policy.

    >**NOTE:** This will completely overwrite the existing RBAC policy of the user. Use with caution.
    """
    # Your logic to override user RBAC policies will come here

You can see how we used the right dependency injection to handle both authentication and authorization. It looks pretty neat, doesn't it? Let's check out some more examples.

@router.get(
    "/dataset/{dataset}/table/{table}",
    dependencies=[Depends(read_table_scope_dependency)],
    response_class=ORJSONResponse,
)
async def read_data(
    dataset: Annotated[DatasetInfoV2, Depends(dataset_presence_dependency)],
    table: Annotated[str, Depends(table_presence_dependency)],
    filter_query: Annotated[FilterQueryParams, Query()],
    request: Request,
) -> ORJSONResponse:
    """Get data from desired dataset & table"""
    # Code to read table data will come here
    return ORJSONResponse(query_data.to_dicts())

Here you can see we used three dependency injections in two different ways

1. read_table_scope_dependency: It is declared in the router decorator as this dependency is not expected to return anything.

2. dataset_presence_dependency & table_presence_dependency: This is declared as Path parameter. That’s because they return some values that will be used downstream.

You will also note that any user RBAC related dependency is not used here. That’s because the dependencies here are multilevel nested dependencies. A level above dependency will perform user authentication & authorization.

Interact with Databricks Lakehouse

This is like the heart of the Data Sharing system. To provide the ability to read & write data in Lakehouse outside of Databricks Spark ecosystem. I have already discussed the API Endpoint part of it, now let’s dive into core logic.

Phase 1: Databricks Delta Sharing

Databricks has a decent native sharing product - Delta Sharing. Follow the docs on how to create a Share & Recipient (this is not the scope of this article). Once it was done, we had to do the following

• To consume the data in Python we had two options - Spark & Pandas. We obviously chose Pandas because that was the whole point.

• The syntax to read the data on surface looks pretty straight forward

import delta_sharing

# <profile-path>: the location of the credential file.
# <share-name>: the value of share= for the table.
# <schema-name>: the value of schema= for the table.
# <table-name>: the value of name= for the table.

data = delta_sharing.load_as_pandas(
f"<profile-path>#<share-name>.<schema-name>.<table-name>"
)

• But in the practical world, it is never this straightforward. As a good governance policy, we used to maintain a dedicated share for every schema and a dedicated recipient for every share. This allowed us to streamline our governance policies. The goal of API was to make every data accessible regardless of which schema it belongs to. So, we had to store multiple credential files and load them dynamically based on table names.

# We stored all the credential (encrypted) in the database. So based on
# given schema, it was read from database
current_cred_config = await shareCredential.find_one(
    shareCredential.schema_name == schema
)
# getting delta share cred
cred = current_cred_config.share_credential.model_dump()
# NOTE - the token is encrypted, so need to decrypt it first
cred["bearerToken"] = decrypt_data(ENCRYPTION_KEY, cred["bearerToken"])
# WARNING - to create delta sharing client it needs cred file physically present. So deleting it
# as soon as client is created
with open("./config.share", "w") as f:
    json.dump(cred, f)

# NOTE - below is the specific `URL` format that delta sharing client needs
table_url = f"./config.share#{current_cred_config.share_name}.{schema}.{table}"
data = delta_sharing.load_as_pandas(table_url)
logger.info(f"successfully retrieve {schema}.{table} for {user_info.email_id}")
# deleting config
cred_config_clean_up()

Pros

• Readily available sharing capability provided by Databricks.

• No limit to how many catalogs can used to share. No limit on Share & Recipient creation.

Cons

• Databricks often updates the Reader and Writer protocols with new features. This means Delta Sharing needs to catch up, and we must ensure we aren't using unsupported table properties.

• The credential created in Databricks for Recipient do get expired. So we need to frequently rotate them. This can be cumbersome.

• But by far the biggest problem (which was eventually a deal breaker for us) with its pandas driver is no support for offset . They do have support for limit, but without offset, limit have very limited benefits.

  • We could not use pagination.

  • We literally have to load all data in memory which is not practical.

• It may not be a Con but Delta Sharing is designed for just reading data. So we could not add Write data functionality.

Due to all these issues, we had to make architectural changes.

Phase 2: Open source tools

As I stated in Part 1, one of the primary reasons to select Databricks lakehose because it uses open source Delta lake format to store table data. So we started to look for other open-source tools that we can utilize here. We settled on the wonderful Polars library (If you're not already using it, you should definitely start right away! I can't recommend it enough).

It has built-in support to read Delta Lake tables, thanks to the Deltalake library. It also supports Lazy read, which works like Python generators, along with Offset and Limit. Plus, since we were already using Polars, we didn't need to make any big changes to our workflows, which was a huge plus for us!

# Reading Delta table Data with few filter options

import polars as pl

# setting up options for reading delta table using polars
storage_options = {
    "account_name": "<storage_account_name>",
    "account_key": "<storage_account_key>",
}
pyarrow_options = {"parquet_read_options": {"coerce_int96_timestamp_unit": "ms"}} # make sure to set the correct timestamp unit
data = pl.scan_delta(
    source="<path_to_delta_table>", # EG - "abfss://<container>@<storage_account_name>.dfs.core.windows.net/path_to/delta_table"
    pyarrow_options=pyarrow_options,
    storage_options=storage_options,
)

# Creating query plan to get the data based on the query params.
if columns: # NOTE - We used this as query param in fastapi to fetch only required columns
    data = data.select(columns)
if row_count: # NOTE - We used this as query param in fastapi to fetch only required absolute no of rows
    data = data.limit(row_count)
elif page: # NOTE - We used this as query param in fastapi to fetch data based on pagination
    data = data.slice(
        offset=page * env_config.page_size, length=env_config.page_size
    )

# executing query plan
query_data = await data.collect_async(streaming=True)
# NOTE - If not using aysnc then use below code
# query_data = data.collect(streaming=True)

# WARNING - fastapi cannot handle polars dataframe while sending response back, so converting
# it to python dict
json_data =  query_data.to_dicts()

# Writing data using Delta Merge

storage_options = {
    "account_name": "<storage_account>",
    "account_key": "<storage_key>",
}

merge_ops_output = (
    data.write_delta( # NOTE - data is a Polars dataframe
        target="<storage_location>",
        mode="merge",
        storage_options=storage_options,
        delta_merge_options={
            "predicate": "<predicate>",  # condition to determine upsert requirement. e.g. "source.id = target.id"
            "source_alias": "source",
            "target_alias": "target",
        },
    )
    .when_matched_update_all()  # updating all columns for a row already present
    .when_not_matched_insert_all()  # inserting all columns for a new row
    # Delta table supports a lot more merge operation, check them out based 
    # your requirement
    .execute()
)

Pros

• This uses open-source tools.

• Supports proper Offset & Limit which allowed us to develop Pagination.

• It supports all Delta Merge operations out of the box which allowed us to develop Write data functionality.

Cons

• To read or write data, we need to know the storage location of the table in advance. To simplify this for users, who only need to provide the table name, we keep a database with the storage locations of all tables. Additionally, there is a process that regularly updates this database.

• The biggest issue with this approach is again the same issue that we had with the previous approach

  • Reader & Writer protocol version. Even here the version was not updated frequently & we could not use the latest features developed by Databricks.

  • As a temporary fix, we used to maintain two tables - the first one was the original table with all the latest table properties & a copy table stripped with all table properties.

Phase 3: Databricks Connect

They say the third is the charm. This was the case even for us too.

As discussed in the previous two phases we had to struggle with multiple pain points. We managed them for some time but none of them were permanent full-proof solutions. So while maintaining the current architecture, we kept our research engine going. We tried multiple approaches & finally settled on Databricks Connect.

It’s a tool provided by Databricks to run Spark (both SQL & PySpark) jobs on Databricks infrastructure remotely. It includes Databricks SQL Connector to work with Spark SQL & Databricks Connect for Python (based on Spark Connect) to work with PySpark.

Let's explore how to create the Cursor to interact with the Delta table in Lakehouse. Databricks offers a bunch of connection and authentication options, which you can check out here. Since we're using it in the application, we went with OAuth machine-to-machine (M2M) authentication.

from databricks.core.config import Config
from databricks.sdk.core import oauth_service_principal
from databricks.sql.exc import RequestError
from databricks.sql import connect


@asynccontextmanager
async def setup_databricks_cursor(app: FastAPI):
    def _oauth_cred():
        return oauth_service_principal(db_cfg)

    def _get_cursor():
        try:
            return connect(
                server_hostname="<databricks warehouse/cluster host>", # make sure to grab serverless warehouse
                http_path="<databricks http path>", # make sure to grab serverless warehouse
                credentials_provider=_oauth_cred,
            ).cursor()
        except RequestError:
            logger.debug("either token expired or connection with databricks lost, getting new")
            return _get_cursor()

    # Adding cursor to FastAPI app state so that it becomes available to all routers
    app.state.cursor = _get_cursor()
    logger.debug("successfully connected to databricks")
    yield
    app.state.cursor.close()
    logger.debug("successfully closed databricks connection")

app = FastAPI(
    debug=True,
    title="Your title",
    lifespan=setup_databricks_cursor, # NOTE - here we are passing the async context manager created above as a lifespan
)

# For Reading we used Databricks SQL Connector 
import polars as pl
from sqlglot import Dialects, select
from sqlglot.errors import ParseError

try:
    columns = select(*filter_query.column) if filter_query.column else select("*")  # NOTE - We used this as query param in fastapi to fetch only required columns
    query = columns.from_(f"{dataset.catalog}.{dataset.dataset}.{table}") # Full table name. EG - "catalog_name.dataset_name.table_name"
    if filter_query.filter:  # NOTE - We used this as query param in fastapi to filter the data based SQL Where clause
        query = query.where(filter_query.filter)
    if filter_query.order_by: # NOTE - We used this as query param in fastapi to order the data based on SQL Order By clause
        query = query.order_by(filter_query.order_by)
    if filter_query.row_count: # NOTE - We used this as query param in fastapi to limit the number of rows fetched
        query = query.limit(filter_query.row_count)
    elif filter_query.page: # NOTE - We used this as query param in fastapi to enable pagination
        query = query.offset((filter_query.page - 1) * 1000).limit(1000)

    # executing query plan
    query_data = pl.read_database(
        query.sql(dialect=Dialects.DATABRICKS), request.app.state.cursor # SQLGlot converts above query to Databricks SQL dialect
    )
    data = ORJSONResponse(query_data.to_dicts()) # FastAPI as native support for ORJSONResponse, which is faster than JSONResponse
# Error coming from SQLGlot due to incorrect sql query
except ParseError as e:
    # You can add your custom exception handling here

One important point to note for using Databricks Connect is the connection gets destroyed in 10 minutes of inactivity and we have to manually create a new connection again. To solve this we used Python context manager.

# Setting up Databricks connect to execute PySpark code

@contextmanager
def spark_manager(profile: str | Config = "DEFAULT"):
    if isinstance(profile, str):
        cfg = _get_databricks_config(profile)
        # NOTE - It's not necessary to use config. But I am using it to make it more flexible &
        # able to use any profile/environment as needed.
    if isinstance(profile, Config):
        cfg = profile
    spark = DatabricksSession.builder.sdkConfig(cfg).getOrCreate()
    yield spark
    # NOTE - We don't need to close the session as it will be automatically closed in case of 10
    # minutes of inactivity by databricks. Also one more added benefit is that some other workload
    # can use the same session if it's not closed.
    # WARNING - If you still want to close the session, you can uncomment the below line.
    # spark.stop()

db_cfg = Config(
    host=env_config.databricks_host,
    client_id="<databricks service principal client id>",
    client_secret= "<databricks service principal client secret>"
    serverless_compute_id="auto", # To use serverless cluster
)

# This will make sure we'll always gets an connection. If there is already one 
# it will use it. If not present then it will create new & use it. 
with spark_manager(db_cfg) as spark:
        # your PySpark code will come here

Pros

• No limitations on Delta tables with any table properties. In fact, we can even query Views which is not possible in the first two phases.

• Since we are using Databricks SQL, it brings all sorts of SQL capabilities like Select, Offset, Limit, Where, Order By, etc.

• We can use Databricks Serverless Compute, that is shared with other workloads, So it requires no additional infrastructure cost.

• No need to maintain any temporary solution.

Cons

• Vendor lock-in as we have to use tools provided by Databricks.

Conclusion

• Wrapping up, creating a solid enterprise data strategy is all about planning and executing well across different stages.

• Our journey from Databricks Delta Sharing to open-source tools like Polars, and finally to Databricks Connect, shows how important it is to stay flexible and keep improving to tackle challenges like data access and performance.

• By focusing on things like modularity, authentication, and authorization, and embracing open-source solutions, businesses can create a scalable and efficient data strategy that keeps up with what users need.

Introduction

Overview of the importance of Data management in enterprises

Data management is vital for enterprises as it

• Ensures that data is accurate, accessible, and secure, which is essential for informed decision-making.

• Helps organizations streamline operations, improve efficiency, and enhance customer experiences by providing reliable data insights.

• Supports compliance with regulatory requirements and reduces the risk of data breaches by implementing robust security measures.

• Drives innovation and competitive advantage by enabling advanced analytics and data-driven strategies.

Key Features of an Effective Data Sharing Solution in Enterprises

Every organization needs a powerful data sharing solution that is

• Powerful enough to share data of any size.

• Agnostic enough to be used by any tool or programming language.

• Scalable enough to support all users all the time.

• Secure enough to ensure the right people have access to the right data.

Achieving this goal will truly make the solution great. Users will have enough confidence, eventually leading to a good adoption rate.

Robust Data Platform Architecture

Here, we'll cover the key architecture choices we made to build a strong data platform that meets all business needs.

Data storage strategy

The What

These are extremely important steps. You have to make the right choice at the beginning, because if you make the wrong choice then data migration will be a difficult task. Usually, the options to choose from or either Data Warehouse or Data Lakehouse.

We choose to go with Data Lakehouse for the following reasons:

• A Data Lakehouse provides a centralized location to store all types of data.

• It supports not only structured data like tables but also unstructured data such as images, videos, and binaries, as it’s built on top of a Data lake.

• Its modular & open design is highly advantageous. By separating storage from data processing, it allows for the use of various tools to read, write, and process data efficiently.

The Where

Out of all the Data Lakehouse out there, we choose Databricks Lakehouse for the following reasons:

• Databricks Lakehouse supports cloud storage objects for data storage. We already use Azure ADLS gen2, which is natively supported by Databricks Lakehouse.

• It utilizes Apache Spark (PySpark + Spark SQL) for processing and transforming data. While many tools can handle big data, none match the capabilities of Apache Spark.

• It employs the Delta Lake format to store data in tables. Delta Lake introduces ACID properties, a key feature that previously deterred users from moving away from Data Warehouses to Data lakes, which paved the way for Data Lakehouse.

The How

There are many well-defined strategies used in storing data. Since we are using the Databricks Lakehouse platform, we went ahead with Databricks’ Medallion Architecture

• Adopting the Medallion architecture has allowed us to enhance data quality by organizing data into three distinct namespaces, each serving a specific purpose, ensuring they do not interfere with one another.

• The Medallion architecture also improves our ability to meet business needs, as the gold layer is specifically designed to cater to those requirements.

• It also provides a unified data management ability to ease engineering efforts.

Facilitate Data Sharing & Consumption

A robust Enterprise data strategy must have equally robust data sharing & consumption capabilities.

The What

• Data Sharing - The goal is to make data accessible throughout the organization.

• Data Consumption - The goal is to make reading and writing data simple and convenient for authorized users and applications by abstracting all the complexity.

The Where

• Both layers - Data Sharing & Consumption - should act as middleware between the user/tool and the Lakehouse.

• We tried multiple options for hosting it but eventually settled on Kubernetes. We were already using it for other products too. However, hosting the application should be purely based on your requirements (and it's a little out of scope for this article too).

The How

• We took inspiration from how two independent services usually communicate & majorly the answer is APIs.

• In 2024, we had two options to choose from either REST API or GraphQL.

• We selected REST API because GraphQL still has limited support across many tools. Almost all tools and programming languages support REST API, making it universally compatible.

Security Architecture

I cannot stress enough that Security Architecture is as important as other architectures like application, scaling, etc. Treating security as a secondary citizen or an afterthought will come back to haunt you later.

To satisfy all our security needs, we broke down the architecture into parts - Data Security & Access control

Data Security

Effective data security measures are crucial for any organization.

The What

• The goal is to have the ability to store data securely

• The goal is to follow any region specific data localization or compliance policies.

• The goal is to restrict unauthorized use of sensitive data.

The Where

• Databricks Lakehouse lets us save table data externally to cloud storage. We use this feature to store all our tables in Azure ADLS Gen2. This ensures our data is securely stored.

• Azure allows us to select the location of our storage account. So we use a multi-region storage account strategy to meet data localization policies.

The How

• Since we are using Azure ADLS Gen2, we get multi-layered security like authentication, access control, network isolation, data protection, advanced threat protection, and auditing.

• The most crucial is played by Databricks Unity Catalog. It goes hand in hand with the multi-region storage account strategy. Here is how we did it

    Step 1: Configure external locations and storage credentials
    - External locations are defined as a path to cloud storage, 
      combined with a storage credential that can be used to access 
      that location.
    - A storage credential encapsulates a long-term cloud credential 
      that provides access to cloud storage.
  
    Step 2: Configure region specific Unity Catalog
    - From above step, now we use region specific storage account 
      storage credentials to create it's dedicated unity catalog. 
    - Once we have the unity catalog created, we can use Spark (SQL or Pyspark)
      to process the data. 
  
    Step 3: Managing all catalogs
    - This is where Databricks really shines. we can have all different 
      region based unity catalog under same workspace. 
    - This way we can run our spark jobs or ETL pipelines, process the data
      save it back to appropriate location. 
    - This way we comply with all policies.
  

Access Control

The Security architecture isn't just about external protection or following policies; we also need to make sure it's secure on the inside too. This can be done using the Access Control principals.

The What

• The goal is to ensure that only authorized users can access specific data, preventing unauthorized access and potential data breaches.

• The goal is to control who can view or modify data, access control helps maintain the integrity of the data, ensuring that it remains accurate and reliable.

• The goal is to provide a clear audit trail of who accessed or modified data, which is essential for accountability and transparency.

The Where

• Since we are using Databricks Lakehouse with External location, we have to implement two Access Control strategies.

• Governance with Databricks Unity Catalog:

  • It allows us GRANT fine grain control like select, read, write, etc to the various objects like table, view, volume, model, etc.

  • We use User Groups principal instead of direct user principles for granting any access.

• Access Control in Azure ADLS Gen2:

  • All the files associated with tables, views, and models are stored in cloud storage. We made sure that any group will not have higher permission than they have over Databricks Unity Catalog.

  • To keep it simple we usually don’t provide write access to storage accounts associated with higher environments to users. Everything is controlled by the Access Control policy in Databricks.

The How

• Databricks offers a highly adaptable governance policy, which we leverage extensively.

• We implement a Zero Trust Policy in combination with Role-Based Access Control.

• For each Schema, we have established roles such as reader, writer, and owner. Depending on the specific use case, Groups are assigned these roles. Subsequently, the User is added to the appropriate Group.

• When a User no longer requires access, they are removed from the group, eliminating the need for frequent modifications to the grants.

• The Service Principal is crucial, especially in the production environment. We only use it to execute jobs.

Conclusion

• Creating an effective enterprise data strategy is essential for harnessing the full potential of data. Focus on robust data management, secure data sharing, and comprehensive security architecture.

• Implement a well-thought-out data platform architecture, like the Databricks Lakehouse with Medallion Architecture, for efficient data storage and processing.

• Facilitate seamless data sharing and consumption through universally compatible APIs for easy access by authorized users.

• Adopting these strategies helps drive innovation, gain a competitive edge, and make informed decisions based on reliable data insights.

Note: I am asumming you familair with Apache Avro file format, its advantages, its shortcomings, etc.

Tip no 1: Use the correct package

Instead of using the official package from Apache Avro use Fast Avro for Python. Trust me the claims made by the author of fastavro mostly holds true.

Tip no 2: Use of schema

Avro relies on schemas. When Avro data is read, the schema used when writing it is always present. You can read the specification docs to understand more about it in detail. I too found it a bit confusing & keep ever forgetting. So here is the rule of thumb that I follow:

schema = {
    'doc': 'A dummy avro file', # a short description 
    'name': 'dummy', # your supposed to be file name with .avro extension 
    'type': 'record', # type of avro serilazation, there are more (see above docs) but as per me this will do most of the time
    'fields': [ # this defines actual keys & their types
        {'name': 'key1', 'type': 'string'},
        {'name': 'key2', 'type': 'int'},
        {'name': 'key2', 'type': 'boolean'},
    ],
}

Tip no 3: Write correctly

The fastavro default write method for some reason does not use any codec or compression algorithm, which defeats the purpose of using avro. See the below screenshot.

from fastavro import writer, parse_schema, reader

# default codec is None 
with open('dummy.avro', 'wb') as out:
    writer(out, parse_schema(schema), more_rows)

# from the above screenshot its best to use deflate. It also have native supoprt
with open('dummy_deflate.avro', 'wb') as out:
    writer(out, parse_schema(schema), more_rows, codec="deflate")

Tip no 3: Read as a generator

Assuming the file size is huge (that will be the case why you had the need to move from JSON to something like Avro) & fastavro do support lazy loading why not use it. Also, it is straightforward

from fastavro import reader
with open("dummy.avro", "rb") as out:
    f = reader(out) # now f is generator
    # you loop it
    for i in f:
        do_something(i)

1. The Old School way

Using the config to control environments & settings is nothing new. Developers have been using it for ages, so why is it necessary to reinvent the wheel? I am not proposing to reinvent the wheel, but modernize it.

The most popular (& probably the most naïve way too 😕) is the use of file-based configs like JSON, YAML, TOML, etc. Apart from being naïve, they have two significant issues:

• they have to be hardcoded

• they cannot be generated dynamically.

This is a big deal breaker when dealing with multiple environments. If you have multiple sources to populate it (e.g. some secret vault, environment variables, hardcoded values, etc) then I would strongly recommend just dropping the idea of using a file-based config to save hardships in your life 😉.

2. Environment Agnostic Config: Generating config programmatically

To achieve this, I use Pydantic's Setting Management application. Anyone familiar with Pydantic will find this familiar & easy to use.

Usually, you must be using

• Something like a .env file to store all secrets & later read them. Or store secrets/configuration directly as environment variables & then read them.

• Other configurations/settings can be hardcoded and written into json/yaml/toml/ini files.

So with the help of Pydantic's BaseSettings we can combine both. Let's dive in to see it in action.

Pydantic's BaseSettings already have support to read from environment variable, .env file, etc (Follow the original docs to see in more details - Settings management - pydantic). So this way we can have some fields (or class attribute) as hardcoded feils & some populate from environment variables under one common Pydantic model. Let's see an example.

from pydantic import BaseSettings, Feild

class Config(BaseSettings):
    env_variable1: str = Feild(description="some description")
    env_variable: str = Feild(description="some description")
    hard_coded1: str = "some hardcoded value"
    hard_coded2: int = 999

From the above example, the first two fields will be automatically populated from matching environment variables, the next two are the hard coded variables.

Note: Since this is based on Pydantic, you can add all sorts of regular Pydantic validators. See the original docs (above) to see all the possibilities.

3. One Config to rule them all

Now coming to the most important part - How to use one config for all possible environments (e.g. DEV, QA, PROD, etc). Actually, it is quite easy, just create a respective Pydantic BaseSettings model for all environments.

from pydantic import BaseSettings, Feild

class LocalSettings(BaseSettings):
    env_variable1: str = Feild(description="some description")
    env_variable: str = Feild(description="some description")
    hard_coded1: str = "some hardcoded value"
    hard_coded2: int = 999

class DEVSettings(BaseSettings):
    env_variable1: str = Feild(description="some description")
    env_variable: str = Feild(description="some description")
    hard_coded1: str = "some hardcoded value"
    hard_coded2: int = 999

class PRODSettings(BaseSettings):
    env_variable1: str = Feild(description="some description")
    env_variable: str = Feild(description="some description")
    hard_coded1: str = "some hardcoded value"
    hard_coded2: int = 999

Since the underlying environment variables will be different for all environments, so will be the populated fields.

4. How to actually consume the config in code

Now we have a single source of config so moving to next part is to how actually consume it in some code. Again there is nothing novel here. What I do is create a function which takes the underlying environment as a parameter as input & return respective config. Also, I usually store all this in config.py

# logic in config.py
from pydantic import BaseSettings, Feild

class LocalSettings(BaseSettings):
# same as above
class DEVSettings(BaseSettings):
# same as above
class PRODSettings(BaseSettings):
# same as above

def get_config(environment: str):
    match environment:
        case "local":
            config = LocalSettings()
        case "dev":
            config = DEVSettings()
        case "prod":
            config = PRODSettings()
    return config


# in some other part of your code/library
from config import get_config
config = get_config("dev")
some_variable = config.env_variable

# NOTE - you don't need to even hardcode environment parameter. What I do is simply create a environment variable for environment it self & use to in function.
config = get_config(os.environ["environment"])
some_variable = config.env_variable
# Now this makes truly environment agnostic & excat same code will work everywhere.

What if you have multiple scopes of multiple config requirements? The pattern remians the same. Add as many as required configs as Pydantic BaseSettings model & return them.

Let me explain with a simple use case of mine. My application needs to support multiple languages. About 80% of the code is generic but there few logic which are language dependent & changes based on underlying language. So I just define them in their respective language Pydantic model. Let's see an example

from pydantic import BaseSettings, Feild

class EnglishConfig(BaseSettings):
    variable1: str = "some thing"
    variable: list = [1,2,3]
    hard_coded1: dict = {}
    hard_coded2: int = 999

class FrenchConfig(BaseSettings):
    variable1: str = "some thing"
    variable: list = [1,2,3]
    hard_coded1: dict = {}
    hard_coded2: int = 999

class HindiConfig(BaseSettings):
    variable1: str = "some thing"
    variable: list = [1,2,3]
    hard_coded1: dict = {}
    hard_coded2: int = 999

Now exactly similar to above logic for environment settings we can we make language dependant or language specific logic as language agnostic.

5. Bringing all things together

Finally, let me show you how my final config.py looks like

from typing import Any, Optional

from pydantic import BaseSettings

class EnglishConfig(BaseSettings):
    regex_pattern_alphanumeric: Optional[str] = "[^0-9a-z/s]"
    list_of_missing_must_include_words = ["Missing", "Must include"]
    list_of_name_prefixes = ["dr", "mr", "mrs", "jr", "sr"]

class SpanishConfig(BaseSettings):
    regex_pattern_alphanumeric: Optional[str] = "[^0-9a-záéíñóúü/s]"
    list_of_name_prefixes = ["sres", "señora"]
    list_of_missing_must_include_words = ["Falta", "Debe incluir lo siguiente"]

class FrenchConfig(BaseSettings):
    regex_pattern_alphanumeric: Optional[str] = "[^0-9a-z\u00C0-\u017F/s]"
    list_of_missing_must_include_words = ["Termes manquants", "Doit inclure"]
    list_of_name_prefixes = ["m", "madame"]

class LocalEnvironmentSettings(BaseSettings):
    common_root_folder: Optional[str] = "/tmp"
    logging_level: Optional[int] | Optional[tuple] = (10,10,10,)
    status_url: str = "https://some-url-dev.com"
    SOME_SECRET: str 
    CONNECTION_STRING: str

class DevEnvironmentSettings(BaseSettings):
    common_root_folder: Optional[str] = "/tmp"
    logging_level: Optional[int] | Optional[tuple] = (10,10,10,)
    status_url: str = "https://some-url-dev.com"
    SOME_SECRET: str 
    CONNECTION_STRING: str

class QAEnvironmentSettings(BaseSettings):
    common_root_folder: Optional[str] = "/tmp"
    logging_level: Optional[int] | Optional[tuple] = (10,10,20,)
    status_url: str = "https://some-url-qa.com"
    SOME_SECRET: str 
    CONNECTION_STRING: str
class PRODEnvironmentSettings(BaseSettings):
    common_root_folder: Optional[str] = "/tmp"
    logging_level: Optional[int] | Optional[tuple] = (10,10,20,)
    status_url: str = "https://some-url-prod.com"
    SOME_SECRET: str 
    CONNECTION_STRING: str
# NOTE - See how I have changed the `status_url` & `logging_level`for all environments & `regex_pattern_alphanumeric` for all languages.

def get_config(language: str, environment: str):
    # setting language based config
    match language:
        case "en":
            language_config = EnglishConfig()
        case "es":
            language_config = SpanishConfig()
        case "fr":
            language_config = FrenchConfig()
        case _:
            raise ValueError(f"given language: {language} must be either from en, es, pt,")
    # setting environment based config
    match environment:
        case "local":
            environment_settings = LocalEnvironmentSettings()
        case "dev":
            environment_settings = DevEnvironmentSettings()
        case "qa":
            environment_settings = QAEnvironmentSettings()
        case "pro":
            environment_settings = PRODEnvironmentSettings()

    class GlobalConfig(BaseSettings):
                global_language_config = language_config
                global_environment_settings = environment_settings

    return GlobalConfig()

Note: As my other blogs, this idea is not limited just to python but can be used anywhere. I have used python to explain the idea. Few modifications & same approach can be applied anywhere.

I don't think anyone will agree that writing comments in your code are a waste of time and effort. Then Why do most people don't really write good & useful comments? Why do they willingly or unwillingly make their own life & experience difficult?

You all must have seen this meme & then laugh & then move on. This is a very serious problem. Hers is one more,

So why do we fall into such a pitfall even knowing pretty well that there will be a pitfall ahead? Here are some of the reasons that I think are primary contributors:

• The comments don't go hand in hand with the code & look out of the place.
• In an already massive code base, it becomes even more difficult to navigate them.
• The rush to commit the code.
• Some unspoken but reality of human behaviour like, "I wrote the code so beautifully that it is self-explanatory" (Yeah, maybe to yourself but not so obvious to others). OR "If I am just going to write the comments then when I will write the actual code" (writing comments is not an afterthought but you should write them along with the code)
• Finally, some people are just lazy. Nothing can be done about them. As they are already on the path which will surely make their life difficult in future.

So not writing comments is more of a philosophical or even behavioural problem rather than technical or knowledge problem.

2. How to write good comments

There are many good blogs out there (this one is very good) which explains what is good content for comments. My focus is on the practicality & ease of access. Following are the steps that you must follow for writing good comments.

2.1 Philosophical change

This can be either very easy or very difficult to adapt. It's upon you. The ideal way would be to have comments more than the actual lines of code. One more thing that will help is while writing a code think that you are not writing this for yourself but for others.

2.2 Tagging comments

Writing just the comments (even if their contents is good) usually is not that helpful because navigating them becomes difficult. So for this, I have come up with a system of 'Tagging Comments'. It is nothing fancy, you just have to start a comment with its type. Here are the type that I use,

• ANCHOR - Used to indicate a section in your file
• TODO - An item that is awaiting completion, address something, etc in future
• FIXME - An item that requires an immediate bugfix
• NOTE - An important note for a specific code section to fetch the attention of a fellow developer
• REVIEW - An item that requires additional review, very useful during pull requests or merges
• DEPRECATED - An item which is no longer being used & will be removed in future
• WARNING - An warring showing for the following item, if not followed respective bad thing can happen
• SECTION - Used to define a region
• LINK - Used to link to a file that can be opened within the editor (See 'Link Anchors')
• EG - An example of what we should expect in the following item

Let me show some practical examples of how they should be used

# 1. Without comments
with open(file_path) as data_file:
    yield from reader(data_file)

# 2. With comments
# TODO - which encoding to use?
with open(file_path) as data_file:
    # spiting out only unit data lazily
     yield from reader(data_file)

Any linter will give a warning in the above example as no encoding was provided. But I didn't know the answer straight away, so I used a TODO tag to resolve it once I'll know the answer.

# saving the current iteration's payload
# FIXME - Input request is also sending a set in the payload, which it should not as set cannot to
# save as json. Temporary fix is to save it as a text file
DataWriter.str_to_txt(
    str(req_body),
    f"/research/sessions/{uuid}/payload.txt",
)

The above code snippet will still work, but its behaviour is incorrect & it must be changed. That's why I used a FIXME tag. Now how it's different from the TODO tag? So FIXME should be used when you know for certain that the following item will definitely turn into a potential bug & TODO should be used for a wide variety which may not be necessary bugs or code-breaking items.

# to let joblib release all workers gracefully from memory. NOTE - It's only needed here # because the same joblib workers from cleaning ops are used by entity extraction ops & RQ does not need
# joblib's multiprocessing
time.sleep(5)

A NOTE tag is used to bring the attention of fellow developers to the following item. This type of comment has more importance than a regular comment.

SECTION tag should use to group certain business logics that can be put under the common bucket. This becomes very helpful in the case of a pipeline.

# NOTE - Following ETL pipeline will only work with prod configuration
# SECTION - Part A: Extract 
collect_data(source)
clean_data(raw_data)
dump_clean_data(clean_data)

# SECTION - Part B: Transform
read_clean_data(clean_data)
transform_data(data)

# SECTION - Part C: Load
load_data_to_db(transformed_data)

# SECTION - Part D: Post-processing
clean_environment()

Everyone has the habit of commenting out certain business logic or some code snippet. And this is not wrong. There can be some valid reason to still keep the commented code snippet. But this becomes extremely confusing for others as they might not be aware of the reason. Ultimately, this leads to difficulty in maintaining the code. So it's better to add a DEPRECATED along with the reason.

# DEPRECATED - blob storage as a source is not required for now & maybe removed completely in future.
#blob_data_source = [unit_source for unit_source in data if unit_source.is_present()]

For EG tag I follow a couple of rules of thumb,

1. If you think certain logic is not clear, then add an EG comment displaying what can be the potential value will be.
2. For every nested loop I write an example regarding what to expect in the next level.

# EG - "Random123 hello @#" will become "andom hello"
regex_pattern = "[^a-z]"
result = re.sub(regex_pattern, "", text)

for key,val in random_dict.items():
    # EG - another_random_dict[val] = 'some str'
    for val in another_random_dict:
        if isinstance(another_random_dict[val], str):
            some_list.append(another_random_dict[val])

2.3 Navigate & automate comments

One of the reasons that I mentioned above is why people don't want to write comments is difficult in navigating them. If I don't have a solution for this problem then there is no point of this blog 😅. So let me introduce you to an amazing VS Code extension - Comment Anchors which I used very heavily for the above use case.

• It searches for the tag & creates a bookmark. So that using its tree structure you can quickly jump to it. Also, it can act as a one-stop to track all important tags like FIXME, TODO, REVIEW, etc. Go through its docs for more info.

• The extension comes with some default tags/strings to bookmark but with the ability to customize it. The following is the one that I am using (You can refer to their docs on how to customize it according to your specific needs)

"commentAnchors.tags.list": [
    {
        "tag": "ANCHOR",
        "iconColor": "default",
        "highlightColor": "#A8C023",
        "scope": "file"
    },
    {
        "tag": "TODO",
        "iconColor": "blue",
        "highlightColor": "#3ea8ff",
        "scope": "workspace"
    },
    {
        "tag": "FIXME",
        "iconColor": "red",
        "highlightColor": "#F44336",
        "scope": "workspace",
        "isBold": true
    },
    {
        "tag": "NOTE",
        "iconColor": "orange",
        "highlightColor": "#FFB300",
        "scope": "file",
        "styleComment": true
    },
    {
        "tag": "REVIEW",
        "iconColor": "green",
        "highlightColor": "#64DD17",
        "scope": "workspace"
    },
    {
        "tag": "SECTION",
        "iconColor": "blurple",
        "highlightColor": "#896afc",
        "scope": "workspace",
        "behavior": "region"
    },
    {
        "tag": "LINK",
        "iconColor": "#2ecc71",
        "highlightColor": "#2ecc71",
        "scope": "workspace",
        "behavior": "link"
    },
    {
        "tag": "DEPRECATED",
        "iconColor": "#B22222",
        "highlightColor": "#B22222",
        "scope": "workspace",
        "behavior": "anchor",
        "isBold": true
    },
    {
        "tag": "WARNING",
        "iconColor": "#B22222",
        "highlightColor": "#B22222",
        "scope": "workspace",
        "behavior": "anchor",
        "isBold": true
    },
    {
        "tag": "EG",
        "iconColor": "#00FFFF",
        "highlightColor": "#eb667d",
        "backgroundColor": "rgba(49, 184, 79, 0.2)",
        "borderStyle": "1px solid #23b2ea",
        "borderRadius": 6,
        "scope": "workspace",
        "styleComment": true
    }
],

You can add this to your VS code settings.json file.

3. How to write a good message commit message

• Writing a good commit message is also extremely important. The same pitfalls as above are true here too.
• Just writing something like updated xyz.py or deleted abc.txt or moved some_file.js or bug fix etc is very bad practice. It does not provide any context & becomes difficult to track changes using git blame.
• How to follow common standards across your team? What should be these standards? No need to reinvent the wheel and just use Conventional Commits.
• It is based on excellent Conventional Commits 1.0.0 spec. You can go through the spec (it is definitely a good read).
• It's very easy, intuitive & fun (as it also supports gitmoji 🤩) to use, trust me. Follow its doc to understand more.

Run the following command on your bash (or any alternative like zsh, fish, etc) to set up auto ssh login.

ssh-copy-id -i "<Public key path>" "<user.name@email.com@host_ip_adress>"

Example:

Output:

2. Windows

Run the following commands, in a local PowerShell window replacing user and host name as appropriate to copy your local public key to the SSH host.

$USER_AT_HOST="your-user-name-on-host@hostname"
$PUBKEYPATH="$HOME\.ssh\id_rsa.pub"

$pubKey=(Get-Content "$PUBKEYPATH" | Out-String); ssh "$USER_AT_HOST" "mkdir -p ~/.ssh && chmod 700 ~/.ssh && echo '${pubKey}' >> ~/.ssh/authorized_keys && chmod 600 ~/.ssh/authorized_keys"

Example:

Output:

• There might be two branches with active development and one of the branch needs some specific (updated) file or directory.
• You don't want to merge the complete branch but just need some specific file(s).

Similarly, there might be 'n' no. of scenarios where the common theme is that instead of merging the complete branch, all you need to do is merges specific file(s)/directory(s).

Using a smart trick which I like to call 'Selective checkout' can do the intended job

git checkout destination
git checkout source sub-directory/
git commit -am "Message."
git pull --rebase
git push

Using this git terminal trick now you can actually perform specific file/directory merge without merging the complete branch.

1. Good technical & fundamental language knowledge
2. Understanding software designing
3. How to make the most use of an IDE.

The third point may seem a little odd to many, in fact, it is mostly undervalued. But to become a productive powerhouse 💪, it is far the most important point to master.

I have spent a lot of time customizing my VS code setup and now I feel I have reached a point where I can confidently present it to the world 🌍 as the most complete setup. 🧰

Note: This will be a long blog, but I have divided it into parts so that a reader can easily jump to his/her interest.

Everyone knows how powerful VS code is & supports almost all languages. The most important of all the great features is its support for a lot of languages out of the box. But tweaking the default setup will boost your productivity 🚀 to the next level.

Table of contents

1. Part 1: Python Development
2. Part 2: Git
3. Part 3: Productivity boosters
4. Part 4: Customization

Part 1: Python development

Python (extension) by Microsoft:

• A compulsory requirement for python development.
• Provides language, debug, test, etc support.
• Latest updates have even brought support for the Jupyter notebooks.

Severity: Must

• Pylance (extension) by Microsoft:

  • Works alongside Python extension to provide performant language support.
  • It adds tons of features into bare-metal Python extension like -
    • Docstrings
    • Signature help, with type information
    • Parameter suggestions
    • Code completion
    • Auto-imports (as well as add and remove import code actions)
    • As-you-type reporting of code errors and warnings (diagnostics)
    • Code outline and navigation
    • Type checking mode
    • IntelliCode compatibility

Severity: Must

Visual Studio IntelliCode (extension) by Microsoft:

• It provides AI-assisted development features for Python, TypeScript/JavaScript and Java developers with insights based on understanding your code context combined with machine learning.
• It AI predictions works fairly well and does not try to intrude on your development.

Severity: Must

There are many other extensions that may not fall under the 'must severity' but are still very helpful, lets continue then,

SonarLint (extension) by SonarSource:

• It is a linter or static analysis tool (free and developed by a very respected company in this domain) that lets you fix coding issues before they exist by analyzing the code.
• It can track Bugs 🐛 and Security Vulnerabilities as you write code. But the best part is the documentation that it provides as a Code smell without even the need to commit the code.
• The installation can a little tricky as initially it needs java to run. It by default handles all installation if sufficient write permission is available else you have to do it manually. But trust me when the installation is done this will help you & your team member to stick to all best practices of python coding and also maintain consistent code. 🫂

Severity: Essential

Pylint (native support):

• Pylint is not an extension but a dedicated linter that can be used independently without VS code, but VS code does have its native support, more here.
• It is a Python static code analysis tool that looks for programming errors, helps to enforce a coding standard, sniffs for code smells and offers simple refactoring suggestions.
• It can be installed as simple as pip install pylint. Follow the above official docs for pylint integration with vs code.
  • But if you ask me using its CLI tool is the preferred way. It can be triggered by pylint path/to/dir for analyzing the complete directory & pylint path/to/some_file.py to analyze any specific file.

Severity: Essential

Note: The most optimal Python linter setup in vs code is a combination of Pylance, SonarLint & Pylint. All these three work independently and does not intrude on each other. Their combination provides the perfect linting experience.

Auto code formatting using Black (native support):

• One of the more important point while working in a team and collaborating on a project is writing clean & consistent code. But dividing our focus on logic & writing clean code always harms productivity. To deal with such situations a formatter should be used.
• VS Code support a wide variety of formatter tools (more here) & automates the code formatting. Following are the steps
  • Goto settings --> Extensions --> Python --> Python › Formatting: Provider and select black from drop down menu. Or if you prefer settings.json then simply and this "python.formatting.provider": "black"
  • While you are on your file/document which you want to format, press right click --> Format code. There is even an option to automatically format the current file on save by adding "editor.formatOnSave": true in settings.json.

• I personally use Black due for the following reasons
  • It philosophy is kind of authoritative & will format the code strictly to its set of rules.
  • I believe we already have to make a lot of critical decisions and formating should not be one of them.
  • Also when the whole team uses Black formatter, then the complete codebase will look & feel consistent and clean.

Severity: Essential

Python Docstring Generator (extension) by Nils Werner:

I am assuming that you are well aware of the Docstring.

• Writing a good informative docstring is very important in terms of documentation. But if you or your team decides to follow any formats like (which you should totally do) Numpy, Google,etc. can results in difficulty for writing consistent & complying with a format/style guide.
• Python Docstring Generator can generate a Docstring template that adheres to the selected format based on type hints/type annotation. So that you have don't have to worry about formatting & just focus on writing the required information.
• You can even jump to the next & previous element in the template using shift & shift + tab keyboard keys respectively.

• To change format goto Settings --> Extensions --> Python Docstring Generator configuration --> Select you desire format from the Auto Docstring: Docstring Format drop-down menu.

Severity: Essential

Python Indent (extension) by Kevin Rose:

• If you have ever felt that you keep on messing with the indentation then extension got you covered. Every time you press the Enter key it automatically adds the correct indent.

Severity: Helpful

Python Type Hint (extension) by njqdev:

• Provides type hint auto-completion for Python, with completion items for built-in types, classes and the typing module.

Severity: Helpful

Python Test Explorer for Visual Studio Code (extension) by Little Fox Team:

• This extension allows you to run your Python Unittest, Pytest or Testplan tests with the Test Explorer UI.
• It provides much better and rich information with ample control.

Part 2: Git

GitLens (extension) by Eric Amodio:

• It supercharges the Git capabilities built into Visual Studio Code. It helps you to visualize code authorship at a glance via Git blame annotations and code lens, seamlessly navigate and explore Git repositories, gain valuable insights via powerful comparison commands, and so much more.
• The best part of this extension is that all the default settings that come out of the box work well. It will improve your Git experience 10x. Even Gitlens can be converted to a standalone tool, yeah it's that powerful 💪.
• There are so many awesome features so I would suggest you watch the video

Severity: Must

Git Graph (extension) by mhutchie:

• It generates beautiful, colourful & informative graphs to visualize all your commits history across all branches. This makes it extremely easy to track the project.
• It even creates every single commit a clickable link that can be used to view git diff among others things.
• It also offers a lot more git functionality which I highly suggest to checkout.

Severity: Must

Git History (extension) by Don Jayamanne:

• View and search git log, history along with the graph and details, previous copy of the file.
• Compare branches, commits, files across commits.

Severity: Essential

Info: You must be thinking 🤔 that extension like vs code's builtin git, gitlens, git graph, git history, seems to have some functionality overlapping, which isn't incorrect. But the important thing is all this extension works well without harming others functionality, so it's totally fine to use them together. In fact, when they all are used together, they make vs code the best Git management tool out there.

Conventional Commits (extension) by vivaxy:

• Before using the extension I would suggest you understand the philosophy behind Conventional Commits
• It brings support to Conventional Commits in vs code.

  Note: Explanation about Conventional Commits is out of the scope of this article but you should must go through it before using this extension.

Severity: Essential

Part 3: Productivity boosters

VS Code Workspace (system functionality):

• Workspace can be used to create a sandbox setup specific to the project or environment. Every single setting, extension, config, etc can be customized to cater specifics needs of a project.
• Creating a workspace is very easy just click on on File & select Save Workspace as

• Initially just created workspace will inherit everything that is present in User settings, then you can now start changing anything that you want to & while doing so you will have the option for saving just to the workspace or globally to the user settings.
• I rate this feature of VS code right at the top 🔝. This is what I generally do:
  • For a python project I set a default python path to be used every time
  • I disabled any extensions that are not required in the project
  • Change terminal-specific settings
  • Hell yeah you can even set workspace specific themes, fonts, etc. which I do all the time 😎. I believe having a visual difference helps to differentiate projects.

Severity: Must

Useful keyboard shortcut:

• Follow this excellent blog by Shubham Khatri. I too from time to give visit here to freshen up the commands.

Severity: Essential

Setting sync (system functionality):

• Settings Sync lets you share your Visual Studio Code configurations such as settings, keybindings, and installed extensions across your machines so you are always working with your favourite setup.
• This way you can maintain a similar & familiar setup with your personal, work, or other personal machines. I personally use this extensively.
• But mind you as awesome as it looks, it can also be proved quickly a double-edged sword. Why?
  • There might be few extensions that you just want on specific devices.
  • There might be a few setting's configs that use path
  • Any other settings or keybindings that you want to use only on a specific device.
• To solve this (potential) conflict vs code provides us extremely granular control over what to sync what to not. Following are few of them (which might be most common):
  • Do not want to sync a specific extension: open the extension page --> Click on Do not sync this extension

• Do not want to sync some specific setting: Goto that specific setting (using settings UI) --> Click on gear ⛮ icon --> Sync this setting

Severity: Essential

Thunder Client (extension) by Ranga Vadhineni:

• Thunder Client is a lightweight Rest API Client Extension. It is basically like a Postman inside vs code so that you don't have to leave vs code at all.

Severity: Essential

Path Autocomplete (extension) by Mihai Vilcu:

• Provides path completion. It supports relative, absolute, workspace path auto-completion.

Severity: Helpful

Comment Anchors (extension) by Exodius Studios:

• Writing comments (even more than the code itself 🧑‍💻) is extremely important in long term, so does an efficient way to track & navigate them. Comment Anchors is the best extension to deal with this task. It supports all languages.
• You can place anchors within comments or strings to place bookmarks within the context of your code. Anchors can be used to track TODOs, write notes, create foldable sections, or to build a simple navigation making it easier to navigate your files. Anchors can be viewed for the current file, or throughout the entire workspace, using an easy to use the sidebar.

• It even supports adding a custom anchor. Following are list of anchor that I used

// My custom anchor related code in `settings.json`

    "commentAnchors.tags.list": [
        {
            "tag": "ANCHOR",
            "iconColor": "default",
            "highlightColor": "#A8C023",
            "scope": "file"
        },
        {
            "tag": "TODO",
            "iconColor": "blue",
            "highlightColor": "#3ea8ff",
            "scope": "workspace"
        },
        {
            "tag": "FIXME",
            "iconColor": "red",
            "highlightColor": "#F44336",
            "scope": "workspace",
            "isBold": true
        },
        {
            "tag": "STUB",
            "iconColor": "purple",
            "highlightColor": "#BA68C8",
            "scope": "file"
        },
        {
            "tag": "NOTE",
            "iconColor": "orange",
            "highlightColor": "#FFB300",
            "scope": "file",
            "styleComment": true
        },
        {
            "tag": "REVIEW",
            "iconColor": "green",
            "highlightColor": "#64DD17",
            "scope": "workspace"
        },
        {
            "tag": "SECTION",
            "iconColor": "blurple",
            "highlightColor": "#896afc",
            "scope": "workspace",
            "behavior": "region"
        },
        {
            "tag": "LINK",
            "iconColor": "#2ecc71",
            "highlightColor": "#2ecc71",
            "scope": "workspace",
            "behavior": "link"
        },
        {
            "tag": "DEPRECATED",
            "iconColor": "#B22222",
            "highlightColor": "#B22222",
            "scope": "workspace",
            "behavior": "anchor",
            "isBold": true
        },
        {
            "tag": "WARNING",
            "iconColor": "#B22222",
            "highlightColor": "#B22222",
            "scope": "workspace",
            "behavior": "anchor",
            "isBold": true
        },
        {
            "tag": "EG",
            "iconColor": "#00FFFF",
            "highlightColor": "#31e0ec",
            "backgroundColor": "rgba(49, 184, 79, 0.2)",
            "borderStyle": "1px solid #23b2ea",
            "borderRadius": 6,
            "scope": "workspace",
            "styleComment": true
        },
        {
            "tag": "@:",
            "iconColor": "yellow",
            "highlightColor": "yellow",
            "scope": "workspace",
            "behavior": "anchor"
        }
    ],

Note: It offers much more features, so I highly suggest reading their description. 📖

Bracket Pair Colorizer 2 (extension) by CoenraadS:

• This extension allows matching brackets to be identified with colours.

Severity: Helpful

Code Spell Checker (extension) by Street Side Software:

• This is a very handy extension for someone like me who makes a lot of typos. It not only checks for typos but also provides correct suggestions too.
• It works with 20+ file types (which obviously covers all the popular one). It supports camelCase, PascalCase, snake_case.
• You can even add words to the global level or even workspace level.

Severity: Essential

Error Lens (extension) by Alexander:

• ErrorLens turbo-charges language diagnostic features by making diagnostics stand out more prominently, highlighting the entire line wherever a diagnostic is generated by the language and also prints the message inline.
• This makes debugging & catching errors comparatively easy.

Severity: Essential

footsteps (extension) by Wattenberger:

• Keep your place when jumping between different parts of your code. This is a VSCode extension that will highlight lines as you edit them, fading as you move away. Jump between lines using ctrl+alt+left and ctrl+alt+right.

Severity: Helpful

Zoom Bar (extension) by wraith13

• Can zoom via GUI in the status bar.

Severity: Helpful

Resource Monitor (extension) by mutantdino:

• Display CPU frequency, usage, memory consumption, and battery percentage remaining within the VSCode status bar.

Draw.io Integration (extension) by Henning Dieterichs

• Draw.io inside vs code.

Severity: Helpful

Part 4: Customization

This is the one place where vs code really shines.

Rainglow theme (extension) by Dayle Rees:

• Rainglow is a collection of colour themes & consists of 320+ syntax and UI themes.
• Colour combinations are excellent. All the themes follow similar categories and hierarchies which makes it super easy to pick a new theme.
• Trust me after installing it you won't need any other theme.

Material Icon Theme (extension) by Philipp Kief

• There are many good options for icons theme in vs code but Material Icon Theme covers most grounds and at the same time all the icons are precise and beautifully designed.

  Window Colors (extension) by Stuart Robinson:

• This extension is a bit unique and fun. Automatically adds a unique colour to each window's activityBar and titleBar.
• It works without harming any existing theme extension & works along with it.
• Why this is useful you ask? If you have multiple windows open (like me all the time 😅) then this adds a new colour & makes it super easy to differentiate.

Let's touch on some theory in brief as this blog focuses on practicality

1. Brief theory

class: A class is a user-defined blueprint or prototype from which objects are created. It wraps all the similar methods (ideally by conventions).

methods: A glorified function that is a class member and will always remain bound to a class.

If you wish to understand the theory, then I would suggest you go through

Types of methods available in Python:

1. True method:
   1. instance method
   2. class method
   3. static method
   4. property (not everyone will agree)
2. Method by conventions (this is what something I have derived)
   1. private method
   2. strict private method

2. Practical use case

The beauty of OOPs in any programming language (that has its support obviously 😎) is the flexibility. There will be always more than one way to do any task at hand, but not all are best practice or practical or pythonic.

Note: I will only touch the syntax/theory part briefly with an assumption that the reader knows the theory part but to see their real-world practical use-case.

Let's begin by writing a sample class that will be used across the complete blog

@dataclass
class SampleClass:
    """This is a sample class to explain practical use-case of all methods
    """    
    xyz: int
    abc: int = 4

    def ins_method(self, no: int):
        """Sample instance method
        Args:
            no (int): any int number 
        """        
        print(self.xyz * no)

    @classmethod
    def cls_method(cls, no: int):
        """Sample class method
        Args:
            no (int): any int number 
        """
        print(cls.abc * no)

    @staticmethod
    def stc_method(no: int):
        """Sample static method
        Args:
            no (int): any int number 
        """
        print(SampleClass.xyz * no)

Now we'll see their practical use case one by one.

2.1 instance method

• The plain, simple, regular method with any frills.
• Practically speaking, this type of method is used 90% of the time. In fact, just using the instance method is what you are going to need all the time.
• It has free access to all attributes & even other methods at the same object level. Due to this flexibility it is most widely used.
• From the above SampleClass example ins_method() is the instance method.

@dataclass
class SampleClass:
    xyz: int
    abc: int = 4

    def ins_method(self, no: int):      
        print(self.xyz * no)

s = SampleClass(xyz=10)
s.ins_method(2)

# Output
20

• Some key point to note here:
  • As instance method are bound to an object of the class, so first an object must be created.
  • The same object then can be used anywhere.

Factory function: Before moving to the next part, it is important to understand factory function) as this is one use-case where classmethod and staticmethod have wide practical application. Again here I will be focussing on practical.

2.2 classmethod

• First, get this, classmethod is not a must have/use in Python OOP. Almost always plain instance method will do the job. But there is one use-case where it fits perfectly, i.e. if you want to use the factory function for accessing class attribute.
• Unlike the instance method where we need to create an object of class first & then use the '.' notation to use it, classmethod can be used without creating an object as it is bound to the class itself and not to the object. Sound too technical let's see an example.

@dataclass
class SampleClass:
    xyz: int
    abc: int = 4

    @classmethod
    def cls_method(cls, no: int):
        print(cls.abc * no)

SampleClass(4).cls_method(2)

# Output
8

• As you can see above we have not created any object for SampleClass but directly used '.' notation with 'SampleClass' itself as classmethod can be bound to class directly. If you compare it with the instance method, there we created an object s then the method ins_method() was being accessed as it is bound to s but not to Sampleclass
• So what is the benefit of all this, not a lot but there are a few:
  • The one obvious, use it as a factory function where you don't want to specifically create an object of the class. Why? Maybe you fear the object size will be too large & you only want to use a specific method.
  • It is a way to tell your other fellow team member or other people that this method doesn't depend on the instance variable (value provided by the user) but on the class variable to which the user doesn't have access.

@dataclass
class SampleClass:
    xyz: int #This is instance variable
    abc: int = 4 # This is class variable

    @classmethod
    def cls_method(cls, no: int):
        print(cls.abc * no)

# Traditional way
class SampleClass:
    abc: int = 4 #This is class variable

    def __init__(self, xyz):
        self.xyz = xyz # This is instance variable

• Here cls_method() only have access to abc which is class variable

• Here ins_method() have access to everything.
• A method that needs to use a class variable as well as user input.

@dataclass
class SampleClass:
    xyz: int
    abc: int = 4

    @classmethod
    def cls_method(cls, no: int):
        print(cls.abc * no)

SampleClass(4).cls_method(2) # Here cls_methos is using class variable - 'abc' as well as user input - 'no'

# Output
8

2.3 staticmethod

• Similar to the 'classmethod' even 'staticmethod' is not a must have/use in Python OOP. In fact, its usage is even less than classmethod. Some dev argues that it is totally useless, see here.
• It shares almost all properties from classmethod except access to neither class variable nor to instance variable as it does not use either self or cls.
• Like the classmethod it can be used as a factory function, but should not. Because
  • We already have classmethod for it.
  • classmethod will work perfectly even if it is not using a class variable.
• So where should it be used? There are no true practical use-case of it and can be skipped entirely. But I do use them in some rare occurrence or edge case, When you have some function outside the scope of a class but feel it is tightly related to the class, it can be included as staticmethod as it does not need any class/instance variable.
  • You want to mimic the private method which doesn't need a class/instance variable. But I would advise instead use a private method (more on this later).

2.4 property

• This is a special type of methods, in fact formally it falls under descriptor. Follow this excellent blog if you want to understand the theory part

  https://www.machinelearningplus.com/python-property/

• @property should be used when you want to return an attribute produced by some function/method.

• As the name suggests the property object should always hold some kind of characteristic of the class. Python's builtin pathlib library is one of the best examples.

2.5 private method

• It is a method to which external users do not have access & it is just used internally.
• Python doesn't have a true private method (like in java). But we can implement it using generally agreed conventions.
• The convention is to add '_' as a prefix to the instance method name. By seeing this name convention, the user will understand that this particular should not be touched.

@dataclass
class SampleClass:
    xyz: int
    abc: int = 4

    def ins_method(self, no: int):
        print(self._pvt_method(no) * no)

    def _pvt_method(self, no: int):
        return (self.abc * no)

• But as mentioned earlier that Python doesn't have a true private method, so the user can still use it. Someone who doesn't the name convention might accidentally use it.

  sc = SampleClass(4)
  sc._pvt_method(4)
  
  # Output
  16
  

2.6 strict private method

• You might not find this term anywhere formally. This is somewhat I have coined. As above mentioned limitation of the private method in python, if you want to enforce it we can hack name mangling feature of python by adding '__' or double underscore as a prefix to the instance method name.

@dataclass
class SampleClass:
    xyz: int
    abc: int = 4

    def ins_method(self, no: int):
        return (self.__strict_pvt_method(no) * no)

    def _pvt_method(self, no: int):
        return (self.abc * no)

    def __strict_pvt_method(self, no: int):
        return (self.abc * no)

# Output
sc.ins_method(4)
64
sc._pvt_method(4)
16
sc.__strict_pvt_method(4)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-31-102204654d54> in <module>
----> 1 sc.__strict_pvt_method(4)

AttributeError: 'SampleClass' object has no attribute '__strict_pvt_method'

• Though it strict but not entirely 100% enforced rule as it can be still used by name mangling syntax

sc._SampleClass__strict_pvt_method(4)

# Output
16

• Then why use it? Because it makes it really hard to use to outside the class as it is not that obvious to use.

Note: I will suggest to further read this thread

https://stackoverflow.com/questions/70528/why-are-pythons-private-methods-not-actually-private

1. Deployment testing in brief 💼

Deployment testing is different from system or unit testing. Let's see in brief (as the original intent of this guide is how to integrate it CI/CD pipeline)

• In unit testing, tests are performed to measure the correctness of the system's individual smaller or unit component.
• In contrast, Deployment testing is a testing stage where two or more software units are joined and tested as one entity but after release or deployment.
• Deployment testing in CI/CD pipeline works best and integrates easily if your system is build using a Microservice approach.

2. General mechanics 🧰

Why

• To ensure every individual microservice is working/behaving correctly.
• Buggy code can be caught after release (& ideally before exposing to the public) and trigger an automated (or manual depends on your use case) rollback.

How

• Running a practical (mocking a real-world use case) example against the microservice. The result must be known beforehand and will be used to compare the output from the test.

When

• Deployment testing must be done as the last component of the CD pipeline.
• It should be triggered after the microservice is successfully deployed.

Where

• It is always performed at your infra/deployment layer (eg. Kubernetes)
• It should be the last step of any CI/CD pipeline.

3. Implementation in Azure DevOps pipeline 🚀

• If your microservice exposes an endpoint (which will be in most of the cases) then all you need is to post a request using REST API (or whatever your microservice supports).
• In my case, the complete microservice was packaged and containerized using Docker and deployed to the Kubernetes cluster. I used a simple python script to extract URL, post request using REST API, and compare the output.
• If you have gone through part 2 of this series regarding integration testing, you will notice that a lot of things are common for both. That's true, the core of both these testing is exactly the same, just executing them differs.

Note: I believe the execution part is more important than the logic (as core logic is the same as integration testing). So I will be focussing more on the execution part.

3.1 Test script

• As I mentioned earlier the core of integration testing and deployment testing is the same so even the testing logic is the same.
• The key difference here is, the deployment testing script should be able to perform the test of all the microservices. To put together in context with integration test, deployment test is a compilation of all individual integration test w.r.t to their microservices.
• As the concept of the test script is already explained in detail here, so I am skipping it, though I will show the overall code at the end just to make this blog less cluttered.

3.2 Perform deployment test in the Kubernetes

You may not be using the Kubernetes at all but the idea behind this is platform agnostic and understanding the flow is key, having said that let's move on

1. As all the microservices are Deployment app (in kubernetes world, more here) we will treat the deployment test as a microservice too. Doing so have tons of benefits like:
   • All the networking requirement will be handled by kubernetes. If you maintained multiple staging envs (like dev, qa, prod), we use the kubernetes namespace to perform env specific test.
   • As both client & user (here its microservice & deployment test service) are internal or at the same level in kubernetes, test latency will be much small.
2. The deployment test pod which will be generated from its k8s deployment app will only include the test script.
3. The CD pipeline of every microservice will need a kubectl exec task at the end because all we need to do is run the python script sitting inside the deployment pod.

   kubectl -n <your namespace> exec po/<deployment test pod name> -- python3 deployment_test.py
   

   Note: In the later section will see how all this can be automated using reference/variables.

3.3 Extracting URL/endpoint

Here comes the magic of kubernets as all the networking is handled by itself. There are a lot of options, but we will be using DNS for Services and Pods as both the service as internal.

• All you need to know:

  1. Name of your k8s deployment app.
  2. Namespace where the pod is currently sitting.

• The IP address for URL will be http://<k8s deployment name>.<namespace>.svc.cluster.local/<your remaining endpoint or sub-page>. Let's see an example. Let's say if k8s deployment name is my-deployment & namespace is dev then it will be http://my-deployment.dev.svc.cluster.local/<some-api>

• Directly using IP address is bad practice as it will keep on changing after every new/restart pod. But the above method, kubernetes will handle this for us.

3.4 Integration with CD Job

I will be using the Release pipeline of the Azure DevOps pipeline for CD Job and focus just on the deployment testing task.

• All you need are multiple kubectl task

• Step 1: We need to extract the deployment test pod name to perform the kubectl exec command.
  1. We will use the kubectl get command to extract the current pod name of the given deployment/app. See the Arguments section carefully. This is where we are extracting the name. The argument will be,

     pods -l app=crs-ai-deployment-test -o jsonpath={.items[*].metadata.name}
     

1. Save the output/name in some reference/variable which can be used in a later stage. This can be done using Output variable > Reference name
2. Output format should be always none. This can be done using Advanced > Output format> none

As you can see from the above screenshot, I am using test as a reference which makes the variable name as a test.KubectlOutput.

• Step 2: kubectl exec deployment testing
  1. We will use the kubectl exec command to run the python script. See the argument section carefully. The test.KubectlOutput which was produced in the previous stage is used now.

• Step 3: (extra) print logs if test fail
  1. Similar to integration testing where logs were printed to the console after the test failed for investigation needs can be done even in deployment testing.
  2. We will need a kubectl log command with --since=10m flag

• See the argument section carefully. Here the reference $(pod.KubectlOutput) is the name of the microservice. Its current pod name can be extracted similarly to how we did for the deployment test pod.
• We also need to change the Control option to Only when a previous task failed as this should be only run when the previous task of `kubectl exec' task performing deployment test failed.

So this is how deployment testing can be automated and integrated into the CD job. Let's see some action.

4. Demo

• First, the pod name of the deployment test pod was extracted. Then the testing was performed. As the test failed, next the logs were printed.

5. Deployment test code

This may change based on your requirement. But you can still refer to this for the idea, as always an idea is platform agnostic.

import sys
import json
import logging
import argparse
import requests
from typing import Dict, Tuple

sys.tracebacklimit = 0
logging.basicConfig(level=logging.INFO, format="[%(levelname)s]: %(message)s")
logger = logging.getLogger(__name__)
parser = argparse.ArgumentParser()
parser.add_argument("--namespace", type=str) # Targeted namespace
parser.add_argument("--deployment_name", type=str) # Targeted test as this is compilation of all individual test

def payload_data(deployment_name: str) -> Tuple[str, Dict]:
    """Prepare payload data for Deployment specifics

    Parameters
    ----------
    deployment_name : str
        Name of deployment to perform testing.

    Returns
    -------
    str
        Api name
    Dict
        Sample data to check for deployment testing

    Raises
    ------
    ValueError
        Must be from supported deployment testing: <here you can add all your test name><eg> research-clarity-id-applicability,research-clarity-id-adv-nonadv
    """
 # TODO: Add new deployment name to the list
    supported_deployment = [
        "research-clarity-id-applicability",
        "research-clarity-id-adv-nonadv",
    ]
    if deployment_name not in supported_deployment:
        raise ValueError(
            f"Given deployment is either wrong or not supported.\nIt must be from {', '.join(supported_deployment)} "
        )

    # TODO: Add all new sample data and API here
    elif deployment_name == "research-clarity-id-applicability":
        api_name = "IDA"
        svc = f"crs-id-applicability-api.{args.namespace}.svc.cluster.local"
        path = "./data/sample_data_research-clarity-id-applicability.json"
    elif deployment_name == "research-clarity-id-adv-nonadv":
        api_name = "adverse_nonadverse"
        svc = f"crs-id-adverse.{args.namespace}.svc.cluster.local"
        path = "./data/sample_data_research-clarity-id-adv-nonadv.json"

    with open(path, "r") as file:
        payload = json.load(file)

    return (api_name, svc, payload)

api_name, svc, payload = payload_data(args.deployment_name)
url = f"http://{svc}/api/{api_name}"
logger.info(f"Testing api @ {url}")

header = {"Content-Type": "application/json"}
logger.info(f"Send API request for {args.deployment_name} deployment testing")
response = requests.request("POST", url, headers=header, json=payload)

# TODO: Add all new assert condition here.
if args.deployment_name == "research-clarity-id-applicability":
    response_data = response.json()
    assert (
        response_data["ida_output_path"].split("/")[-1] == "IDA.ndjson"
    ), "Not Received expected output, test is failed"
elif args.deployment_name == "research-clarity-id-adv-nonadv":
    response_data = str(response.content).replace("'", "")
    assert (
        response_data.split("/")[-1] == "classifiation_output.ndjson"
    ), "Not Received expected output, test is failed"

logger.info("Deployment test passed successfully !!!")

1. Integration testing in brief 💼

Integration testing is different from system or unit testing. Let's see in brief (as the original intent of this guide is how to integrate it CI/CD pipeline)

• In unit testing, tests are performed to measure the correctness of individual smaller or unit component of the system.
• In contrast, Integration testing is a testing stage where two or more software units are joined and tested as one entity.
• Integration testing in CI/CD pipeline works best and integrates easily if your system is build using a Microservice approach.

2. General mechanics

Why

• To ensure every individual microservice is working/behaving correctly.
• Buggy code will be restricted during the CI pipeline only & will never get deployed.

How

• Running a practical (mocking a real-world use case) example against the microservice. The result must be known beforehand and will be used to compare the output from the test.

When

• Integration testing must be done during the CI pipeline.
• It should be triggered after the microservice is successfully build.

Where

• It is always performed at Remote Repository (eg Github, Gitlab).
• It should be the second step of any CI/CD pipeline.

3. Implementation in Azure DevOps pipeline 🚀

• If your microservice exposes an endpoint (which will be in most of the cases) then all you need is to post a request using REST API (or whatever your microservice supports).
• In my case, the complete microservice was packaged and containerized using Docker. I used a simple python script to extract URL, post request using REST API, and compare the output.

Let's go first through the testing script and then through the CI pipeline

3.1 Test script

1. The core of the script is to read some sample data send a post request and finally compare result. The bare minimum would be

with open('./tests/sample_data.json', 'r') as f:
    payload = json.load(f)

header = {"Content-Type": "application/json"}

response = requests.request('POST', url, headers=header, json=payload)
response_data = response.json()

assert response_data['ida_output_path'].split('/')[-1] =='IDA.ndjson', 'Not Received expected output, test is failed.' # You may use some other method to compare 😁

• This will only work in an ideal scenario which will be not possible 99% of the time & in fact completely defeats our original purpose.
• We need to make it more suitable & versatile for this use case. It can be done by adding two more components,
  1. A try-catch block.
  2. console logging at the time of failure to investigate it.

try:
    response = requests.request('POST', url, headers=header, json=payload)
    response_data = response.json()
except:
    logging.error('An error has occurred. Refer logs to locate error.')
    os.system("docker logs test_api > output.log")
    time.sleep(3)
    with open('./output.log', 'r') as log:
        print(log.read())

try:    
    assert response_data['ida_output_path'].split('/')[-1] =='IDA.ndjson', 'Not Received expected output, test is failed.'
except (AssertionError,KeyError) as e:
    os.system("docker logs test_api > output.log")
    time.sleep(3)
    with open('./output.log', 'r') as log:
        print(log.read())

• Some pointers from the above code block
  1. I have added post request and result comparison in the try-catch block so that any error will not stop the code.
  2. As I have mentioned earlier I am performing the test inside the docker image so all the logs I am extracting from it (as shown on line no 6, 15)
  3. Printing logs at the time of failure is very important for investigating the issue. The aim of Integration testing is not just to restrict a buggy release but also to help to investigate it.

3.2 Extracting URL/endpoint

• This may differ quite a lot based on your use case. This is how I do it.
  1. I first start/run the freshly built docker container on the worker node of CI pipeline
  2. Then extract the IP address where the container is running.
  3. Use this IP as the endpoint to post request over REST API.

subprocess.call(['docker', 'run', '-d', '-p' ,'80:80','--name', 'test_api', args.image_name])
ip = subprocess.getoutput("docker inspect --format='{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}' test_api")
url = f'http://{ip}/api/IDA'
logger.info(f"Testing api @ {url}")

3.2 Integration with CI Job

• If you follow a similar flow then all you need is a task to run the python script. It can be as simple as,

- task: PythonScript@0
  displayName: Integration testing
  inputs:
    scriptSource: 'filePath'
    scriptPath: '$(System.DefaultWorkingDirectory)/tests/integration_testing.py'
    arguments: '--image_name your.repo.io/ida:$(Build.BuildNumber)'

• There are two more important points for implementing Integration testing
  1. Placement: It should be performed just after Docker build and before Docker push
  2. If the test fails the script must stop the pipeline to proceed further. This can be achieved by using something like os.sys.exit()

4. Working/Demo

Img: Case- Passing of Integration test

Img: Case- Failing of Integration test.

• As you can see no further task were executed after the failure of the integration test and eventually any subsequently connected CD tasks.

5. Complete code

1. integration_testing.py

import os
import json
import time
import logging
import argparse
import requests
import subprocess


logging.basicConfig(level=logging.INFO, format='[%(levelname)s]: %(message)s')
logger = logging.getLogger(__name__)
parser = argparse.ArgumentParser()
parser.add_argument('--image_name', type=str)

args = parser.parse_args()

# Run docker container at port 80
subprocess.call(['docker', 'run', '-d', '-p' ,'80:80','--name', 'test_api', args.image_name])
ip = subprocess.getoutput("docker inspect --format='{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}' test_api")
url = f'http://{ip}/api/IDA'
logger.info(f"Testing api @ {url}")
with open('./tests/sample_data.json', 'r') as f:
    payload = json.load(f)

header = {"Content-Type": "application/json"}
logger.info('Waiting for 30 sec to let docker container start')
time.sleep(30)
logger.info("Send API request for integration testing")
try:
    response = requests.request('POST', url, headers=header, json=payload)
    response_data = response.json()
except:
    logging.error('An error has occurred. Refer logs to locate error.')
    os.system("docker logs test_api > output.log")
    time.sleep(3)
    with open('./output.log', 'r') as log:
        print(log.read())
    os.sys.exit('Task terminated')

try:    
    assert response_data['ida_output_path'].split('/')[-1] =='IDA.ndjson', 'Not Received expected output, test is failed.'
except (AssertionError,KeyError) as e:
    os.system("docker logs test_api > output.log")
    time.sleep(3)
    with open('./output.log', 'r') as log:
        print(log.read())
    logging.error(f'Problem with {e} Refer logs to locate error')
    logging.error(f"Response from API: {response_data}")
    os.sys.exit('Task terminated')

logger.info("Integration test passed successfully moving to next task")

1. CI pipeline

trigger:
  branches:
    include:
    - develop
  paths:
    exclude:
    - Dockerfile_base
    - requirements.txt

pool:
  vmImage: 'ubuntu-latest'

steps:

- checkout: self
  clean: true
  fetchDepth: 1

- task: UsePythonVersion@0
  inputs:
    versionSpec: '3.x'
    addToPath: true
    architecture: 'x64'

- task: CmdLine@2
  displayName: Install python package
  inputs:
    script: 'python3 -m pip install requests azure-devops'

- task: PythonScript@0
  displayName: Check build pipeline status
  inputs:
    scriptSource: 'filePath'
    scriptPath: '$(System.DefaultWorkingDirectory)/tests/base_pipeline_status.py'
    arguments: '--personal_access_token $(PERSONALACCESSTOKEN) --repo_id b978e55f-bf80-466c-86c8-fc0dfe909b2c --pipeline_def_id 179'

- task: Docker@0
  displayName: 'Docker Build Image'
  inputs:
    azureSubscription: 'Your Subscripton'
    azureContainerRegistry: 'Your container registery'
    dockerFile: Dockerfile
    buildArguments: |
      ARG_STORAGEACCOUNTNAME=$(STORAGEACCOUNTNAME)
      ARG_CONTAINERNAME=$(CONTAINERNAME)
      ARG_STORAGEACCOUNTKEY=$(STORAGEACCOUNTKEY)
      ARG_MAXWORKERS=$(MAXWORKERS)
    imageName: 'ida:$(Build.BuildNumber)'

- task: PythonScript@0
  displayName: Integration testing
  inputs:
    scriptSource: 'filePath'
    scriptPath: '$(System.DefaultWorkingDirectory)/tests/integration_testing.py'
    arguments: '--image_name your.repo.io/ida:$(Build.BuildNumber)'

- task: Docker@0
  displayName: 'Push image to ACR'
  inputs:
    azureSubscription: 'Your Subscripton'
    azureContainerRegistry: 'Your container registery'
    action: 'Push an image'
    imageName: 'ida:$(Build.BuildNumber)'

- task: PublishBuildArtifacts@1
  inputs:
    PathtoPublish: '$(Build.SourcesDirectory)/kube'
    ArtifactName: 'drop'
    publishLocation: 'Container'

TDD is a lot more than just the Unit testing. Lets us see how we can design our CI/CD pipeline to adapt to the TDD approach.

Info: 💡💡💡

I am using the Azure DevOps pipeline as CI/CD tool. But I believe you will be able to implement it in other CI/CD tool too. Understanding the mechanism is key.

Let's break down the pipeline into three component and see how to perform a test there.

1. Pull request testing in brief 💼

PR testing is different from the system or unit testing. Let's see in brief (as the original intent of this guide is how to integrate it CI/CD pipeline)

• In unit testing, tests are performed to measure the correctness of the system's individual smaller or unit component.
• In contrast, PR testing is a testing stage where the functional test is performed on the complete scope branch which has to be merged.
• PR testing in CI/CD pipeline works best and integrates easily if your system is build using a Microservice approach.

2. General mechanics 🧰

Why

• To ensure the functionality code is working as expected.
• It helps to restrict the merging of any unintended commits.
• Ensures nothing breaks over collaboration.

How

• The core of running a PR test is the same as running a Unit test.
• All the tests which is used here are identical to Unit testing with a key difference of:
  1. Unit test are/should be performed locally whereas PR testing is performed over the remote.
  2. Code coverage of PR testing must be always more than the unit test because the intention of the unit test is to just check individual changes locally, whereas PR testing intention is to check for the complete branch which may more than one collaborative effort.
  3. It may be possible that you may not need a 100% passing of the Unit test (though this is not ideal) but PR testing must be 100% passing.

When

• PR testing should be performed at the time of a PR completion or merge activity.

Where

• It is always performed at Remote Repository (eg Github, Gitlab).
• It should be the first step of any CI/CD pipeline.

3. Implementation it in the Azure DevOps pipeline 🚀

• I am using python and pytest here. I am assuming that you will already have all the unit tests written and follow all conventions for pytest automatic test discovery.
• All you need to run is a terminal command

- task: CmdLine@2
  inputs:
    script: |
      # If you have setup.py
      python3 -m pip install .
      pytest -vv
    workingDirectory: '$(System.DefaultWorkingDirectory)'

• But the critical part here is How to make it automated? After all, this is part of CI/CD and the whole point of CI/CD is automation.
• This is a two-step process in Azure DevOps:
  1. Create a new pipeline dedicated to running the unit tests. The trigger should be the same as its parent build or CI pipeline.

trigger:
- develop

pool:
  vmImage: 'ubuntu-latest'

steps:

- task: UsePythonVersion@0
  inputs:
    versionSpec: '3.x'
    addToPath: true
    architecture: 'x64'

- task: CmdLine@2
  inputs:
    script: |
       # If you have setup.py
      python3 -m pip install .
      pytest -vv
    workingDirectory: '$(System.DefaultWorkingDirectory)'

Tip: 🤝🤝🤝

From above you can see I am using trigger as develop branch as this PR testing is intended for develop branch.

1. Branch Policies:
   • For enabling PR triggers we need to branch policies in Azure DevOps
   • Goto Repo > Branch > Branch Policies > Build Validation

Img: Branch Policies Img: Build Validation

• We have to add a build validation policy for the PR trigger. Out of all two settings is key,
  1. Trigger: Must be automatic
  2. Policy requirement: Required

Img: Build Validation Policy

• See in the Build pipeline section, I am pointing to the PR testing pipeline that I showed earlier.

Img: How the PR testing works (see the marked item)

It is vital to understand the past responsibility to adapt to the current expectation of the industry.

• The AI space was still relatively new (though not in academics) and many companies, startups were analyzing its application and valid use-case.
• The research was the primary focus. The caveat here was that this research many times was not directly in line with the core of the organization. So initially not much credibility was expected.
• Generally, companies used to blend the roles of a Data Scientist with a Data analyst or Data engineer. Again, due to the vagueness of AI enterprise application.
• Individuals also had a kind of similar dilemma. A lot of their research or work was not directly in line and practically not viable to be served as a product.

2. The current outlook

The democratization of AI has seen remarkable developments from businesses and startups. Let us try to understand it,

• The industry now distinguishes the role of a Data Scientist, Machine Learning Engineer, Data Analyst, Data engineer, even MLops engineer.
• Businesses no longer allow research in the wild, as they know what use-case exactly they are tapping in. A clear mindset & similar discrete approach from an individual is also required.
• Every Research or POC must have a tangible and servable product

3. The thorough dissection of all the Roles

If we have to pick one area where the Businesses have excelled in AI space, it is undoubtedly the clear expectation from all varieties of the Roles, which are in a nutshell:

1. Data Scientist: A Data Scientist is a person who (generally from a stats/maths background) uses a variety of means including AI to extract valuable information from data.

   • A fundamental difference between Data Analyst & Data scientist is- the former generally rely on domain knowledge and manual old school methods to make sense of data on a small to medium scale, whereas, the latter is responsible for collecting, analyzing and interpreting data on a larger scale using wider means of tools like AI, SQL, old school manual ways, etc.,
   • Domain knowledge is not a must but having is helpful.
   • The primary job is to maintain and extract business contributing insights from data & not to develop the software or product.
   • A Statistician or a Mathematician can become a good Data Scientist.

2. Machine Learning Engineer: A niche software engineer who develops a product or service based on AI.

   • An ML engineer needs to have all the expertise of traditional software engineering along with knowledge of AI because he/she is eventually going to build software with AI at its core.
   • The primary job is not to extract data but to develop an AI tool that can perform the same job.
   • A developer with good knowledge of machine learning/deep learning as well as software engineering can become a good Machine learning engineer.

3. Machine Learning Operation Engineer: A niche software engineer who maintains and automates the pipeline which is used by the ML system.

   • Relatively new field inspired by DevOps. Though different from traditional DevOps roles.
   • Unlike traditional software engineering, development for any product/software/service based on AI doesn't stop at the completion of the building of software. It has to be updated regularly with new data, based on the Data-Drift.
   • The primary job includes all traditional DevOps work as well as maintaining/automating pipeline and Data-Drift
   • A developer with good knowledge of machine learning/deep learning, software engineering & cloud technologies can become a good MlOps engineer.

4. Data Engineer: A niche software engineer who develops a pipeline to serve all data needs using a variety of tools (generally cloud-based)

   • The data engineer needs to have expertise with major cloud-based data platform, batch or stream processing, big data platform (depends on business requirement), database (though not like a Database administrator).
   • Have to work in line with Data Governance policies.
   • Primary job is to design, implements and maintain the Data pipeline.
   • A developer with good knowledge of cloud technologies, data platform, processing needs can become a good Data engineer.

For a new seeker or someone who is aiming to advance in his or her career, all these roles and expectations must be well understood. Given that companies are clearly distinguishing this role, it is expected that this will also be the case for individuals. A vague mindset is totally useless.