Python Virtual Environment Tasks in dbt pipelines¶

# import the necessary libraries what are not available on airflow workers
import pandas as pd
import numpy as np
from sqlalchemy import create_engine


# HACK: we need to define this function so that dbt can parse the model correctly
# here we need to define all refs that are used in the model
def model(dbt, session):
    """
    Here you can reference any dbt model or source to build the dependency graph (it's not required).
    Although dependencies can be specified, your code must perform actual reading of data.
    """
    # dbt.ref('some_model')
    # dbt.source('some_source', 'some_table')

    return pd.DataFrame()  # it's required to return DataFrame in dbt python models


def ingest_venv():
    # do some work here
    print("Ingesting data")
    df = pd.DataFrame(np.random.randint(0, 100, size=(100, 4)), columns=list('ABCD'))
    print(df.head())
    # save the data to the postgres
    engine = create_engine('postgresql://airflow:airflow@localhost:5432/airflow')
    df.to_sql('random_data', con=engine, if_exists='replace', index=False)


# run the function
ingest_venv()