Automate stock market data updates with CrateDB and Apache Airflow

Watch this tutorial on YouTube: Automating stock data with Airflow and CrateDB.

This guide shows how to automate collecting and storing stock market data for S&P 500 companies.

Quick overview

Let’s have a quick overview of what you’ll do:

Goal:

Update stock market data regularly.

Approach:

Define tasks to download, prepare, and store data; orchestrate them with Airflow.

Steps:

Start CrateDB and create a table; create an Airflow project and set the CrateDB connection; implement the DAG in Python; schedule it.

Setup

Set up on macOS. Ensure Homebrew is installed and Docker Desktop is running.

Run CrateDB and create a table to store data

First, run CrateDB with Docker. With Docker Desktop running, copy the command from the CrateDB installation page and run it:

docker run --publish=4200:4200 --publish=5432:5432 --env CRATE_HEAP_SIZE=1g crate:latest '-Cdiscovery.type=single-node'

With CrateDB running, you can now access the CrateDB Admin UI by going to your browser and typing localhost:4200.

Create a table to store financial data. Focus on the adjusted close value (“adjusted_close”) per ticker per day. Use a composite primary key on (closing_date, ticker):

CREATE TABLE IF NOT EXISTS doc.sp500 (
   closing_date TIMESTAMP,
   ticker TEXT,
   adjusted_close DOUBLE PRECISION,
   primary key (closing_date, ticker)
);

We are done with the Admin UI for now. Let’s return to the terminal to install Astronomer.

Install Astronomer CLI and initialize the project

In a new terminal tab, install Astronomer:

brew install astro

Create a new directory for the Airflow project:

mkdir astro-project && cd astro-project

Initialize the project:

astro dev init

Now you have the skeleton of your Airflow project, which looks like this:

├── dags # directory containing all DAGs ├── include # additional files which are used in DAGs ├── .astro # project settings ├── Dockerfile # runtime overrides for Astronomer Docker image ├── packages.txt # specification of OS-level packages ├── plugins # custom or community Airflow plugins ├── setup # additional setup-related scripts/database schemas └── requirements.txt # specification of Python packages

By default, PostgreSQL listens on 5432 and the web server on 8080. If these ports are in use, change them as shown below.

Last adjustments

There are now three things you have to adjust before running Airflow:

  • Add your CrateDB credentials to the .env file. Open the file in a text editor, and add the following line, which takes the default credentials for CrateDB, with user = crate, and password = null. (note: my internal port for running CrateDB in Docker is 5433, which I use here. If using the standard Docker command with 5432, here it should also be 5432).

    # For local development only; do not commit real credentials
AIRFLOW_CONN_CRATEDB_CONNECTION=postgresql://crate:crate@host.docker.internal:5433/?sslmode=disable

  • If the default ports are unavailable, you can change them to free ports. Just open the .astro/config.yaml file in a text editor and update the web server port to 8081 (instead of default 8080) and Postgres port to 5435 (instead of the default 5432), like so:

    project:
   name: astro-project
webserver:
   port: 8081
postgres:
  port: 5435

Start Airflow

Now you are done with the last adjustments, head back to your terminal and run this command to start Airflow: astro dev start You can now access Airflow in your browser at http://localhost:8081.

Write the DAG

In Airflow, define tasks as nodes in a DAG—a Directed Acyclic Graph. That means you set the tasks to run one after the other without cycles to avoid deadlocks. A task (or node) does not stand by itself: it depends on other tasks, and other tasks depend on it. These dependencies are the edges of the Graph and make up the DAG structure by connecting the tasks. You bring the DAG to life by writing the tasks in Python with the help of Airflow operators and Python modules. Now you’ve learned enough to start building your DAG step-by-step!

Create astro-project/dags/financial_dag.py. The DAG file has the following structure:

  • Import operators and python modules

  • Declare functions

  • Set DAG and its tasks

Import operators and modules

Import the operator used in this guide, SQLExecuteQueryOperator, and the decorator to define the DAG and its tasks. You will also import the datetime, pendulum modules to set up your schedule and the yfinance, pandas, and json modules to download and manipulate the financial data later.

import datetime
import math
import json
import logging
import pendulum
import yfinance as yf
import pandas as pd
from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
from airflow.decorators import dag, task

Don’t forget to add these modules to the requirements.txt file inside your project like so:

apache-airflow-providers-postgres>=5.3.1
apache-airflow-providers-common-sql>=1.3.1
apache-airflow[pandas]
yfinance==0.1.87

Declare tasks

The next step is to declare the necessary tasks to download, prepare and insert data.

Download task

Let’s first write a function to download data from yfinance; I will call it download_yfinance_data. You can use ds for today’s date or get yesterday’s date with airflow.macros.ds_add(ds, -1). You start by listing tickers from stocks of interest into a tickers variable. You then pass this list and the start date as arguments to the yf.download function and store the result in a data variable. data is a pandas data frame with various values for each stock, such as high/low, volume, dividends, and so on. Today, I will focus on the adjusted close value, so I filter data using the Adj Close key. Moreover, I return the data as a JSON object (instead of a data frame) because it works better with XCom, which is Airflow’s mechanism to talk between tasks. Finally, you set this function as an Airflow task using the @task decorator and give it an execution timeout.

@task(execution_timeout=datetime.timedelta(minutes=3))
def download_yfinance_data(ds=None):
    """Downloads Adjusted Close data from S&P 500 companies"""

    # list of stocks of interest
    tickers = ["AAPL", "AMZN", "META", "TSLA"]
    data = yf.download(tickers, start=ds)["Adj Close"]
    return data.to_json()

Prepare data task

Next in our code comes a prepare_data task, making the insert task more manageable later. In this function, you get the data in JSON as a parameter and transform it into a data frame for easier manipulation. Then, you take the closing_date, ticker, and adj_close columns from the sp500 table you created and make them keys in a dictionary. You turn each data frame row into a dictionary with these keys. Finally, you add each of these dictionaries into a values_dict list, and done! Mark it as a task with the @task decorator, and give an execution timeout. Now the data is ready for the last task: insert it into CrateDB.

@task(execution_timeout=datetime.timedelta(minutes=3))
def prepare_data(string_data):
    """Creates a list of dictionaries with clean data values"""

    # transforming to dataframe for easier manipulation
    df = pd.DataFrame.from_dict(json.loads(string_data), orient="index")

    values_dict = []
    for col, closing_date in enumerate(df.columns):
        for row, ticker in enumerate(df.index):
            adj_close = df.iloc[row, col]

            if not (adj_close is None or math.isnan(adj_close)):
                values_dict.append(
                    {
                        "closing_date": closing_date,
                        "ticker": ticker,
                        "adj_close": adj_close,
                    }
                )
            else:
                logging.info(
                    "Skipping %s for %s, invalid adj_close (%s)",
                    ticker,
                    closing_date,
                    adj_close,
                )

    return values_dict

Insert data task

So you have a list of dictionaries, each having values to import into CrateDB. For each of these values, you want to execute the INSERT statement, and that’s when the expand method comes in handy. What Airflow does here is similar to defining the task in for loop: it creates n copies of the task, one for each input.

You use the SQLExecuteQueryOperator to execute SQL statements against CrateDB using the cratedb_connection you defined in the beginning. Your final task looks like this:

SQLExecuteQueryOperator.partial(
    task_id="insert_data_task",
    conn_id="cratedb_connection",
    sql="""
        INSERT INTO doc.sp500 (closing_date, ticker, adjusted_close)
        VALUES (%(closing_date)s, %(ticker)s, %(adj_close)s)
        ON CONFLICT (closing_date, ticker) DO UPDATE SET adjusted_close = excluded.adjusted_close
        """,
).expand(parameters=prepared_data)

Main execution method

Finally, it’s time to wrap everything up in your main execution method, financial_data_import. Here, you start by storing the downloaded data from download_yfinance_data into a yfinance_data variable and then storing in prepared_data the results from prepare_data. Then, the SQLExecuteQueryOperator is called to perform the INSERT.

def financial_data_import():
    yfinance_data = download_yfinance_data()

    prepared_data = prepare_data(yfinance_data)

    SQLExecuteQueryOperator.partial(
        task_id="insert_data_task",
        conn_id="cratedb_connection",
        sql="""
            INSERT INTO doc.sp500 (closing_date, ticker, adjusted_close)
            VALUES (%(closing_date)s, %(ticker)s, %(adj_close)s)
            ON CONFLICT (closing_date, ticker) DO UPDATE SET adjusted_close = excluded.adjusted_close
            """,
    ).expand(parameters=prepared_data)

Set DAG and its tasks

Now that you have your main execution method, it’s time to put the DAG together. You start with the @dag decorator and define the schedule parameter; in my case, I’ll make it run daily. You can also give it a start_date for the first DAG run, for example, from 2022-12-01. If the start date is before the current date, by default, Airflow will attempt to backfill all missed DAG runs. If you don’t want that to happen, set the catchup parameter to false.

@dag(
    start_date=pendulum.datetime(2022, 1, 10, tz="UTC"),
    schedule="@daily",
    catchup=False,
)

Your DAG structure will look like this: you first add the operators and modules, then the tasks. Then comes the DAG definition followed by the primary execution method: in this case, financial_data_import. The last line calls the main method, and your Airflow DAG is ready!

import datetime
import math
import json
import logging
import pendulum
import yfinance as yf
import pandas as pd
from airflow.providers.common.sql.operators.sql import SQLExecuteQueryOperator
from airflow.decorators import dag, task

@task(execution_timeout=datetime.timedelta(minutes=3))
def download_yfinance_data(ds=None):

@task(execution_timeout=datetime.timedelta(minutes=3))
def prepare_data(string_data):


@dag(
    start_date=pendulum.datetime(2022, 1, 10, tz="UTC"),
    schedule="@daily",
    catchup=False,
)
def financial_data_import():
    yfinance_data = download_yfinance_data()

    prepared_data = prepare_data(yfinance_data)

    SQLExecuteQueryOperator.partial(
        task_id="insert_data_task",
        conn_id="cratedb_connection",
        sql="""
            INSERT INTO doc.sp500 (closing_date, ticker, adjusted_close)
            VALUES (%(closing_date)s, %(ticker)s, %(adj_close)s)
            ON CONFLICT (closing_date, ticker) DO UPDATE SET adjusted_close = excluded.adjusted_close
            """,
    ).expand(parameters=prepared_data)


financial_data_import()

Execute DAG in Airflow UI

Now that your DAG code is ready, you can interact with it from the Airflow UI. Navigate to http://localhost:8081 and enter the default credentials (user = admin, password = admin) to access the Airflow UI. You should now see the DAG you just created under the DAGs tab. Airflow UI Next to the DAG’s name, you find a toggle button to pause/unpause the DAG’s execution. And if you click on the DAG’s name, you get redirected to a page with different views for your DAG. For example, clicking on Graph will show you the tasks you have recreated; you can check out your code on the Code tab, and the Grid gives you information about your last runs and access to the logs.

A play button on the top right triggers the DAG manually. So let’s click on it to test our DAG. Trigger DAG

On the Grid view, you can click on these squares to check out the execution details. Grid view

Check out the data in the Admin UI

Now that your task is executed let’s check the data in CrateDB! In the Admin UI, select your table in the tables tab and click Query Table. You should now see some of the records you just imported! From now on, you can easily access your financial data in CrateDB and use it as you like. Admin UI