Spaces:
Runtime error
Runtime error
File size: 1,054 Bytes
6defa3d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
from zenml import step, pipeline
import psycopg2
from typing import Tuple
from typing_extensions import Annotated
import logging
import os
from dotenv import load_dotenv
import pandas as pd
from src.data_retrieving import Baseline
load_dotenv()
sql_folder_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "sql_files")
@step(enable_cache=False)
def retrieve_data() -> Tuple[
Annotated[pd.DataFrame, "df_movie"],
Annotated[pd.DataFrame, "df_user"],
Annotated[pd.DataFrame, "df_rating"],
]:
conn = psycopg2.connect(
dbname=os.environ.get('DB_NAME'),
user=os.environ.get('DB_USER'),
password=os.environ.get('DB_PASSWORD'),
host=os.environ.get('DB_HOST'),
port=os.environ.get("DB_PORT")
)
data_retrieve_class = Baseline()
df_movie, df_user, df_rating = data_retrieve_class.retrieve_data_sql(conn)
conn.close()
return df_movie, df_user, df_rating
@pipeline
def postgres_pipeline():
connect_to_postgres()
if __name__ == "__main__":
postgres_pipeline() |