Run local ruff formatting and set continue-on-error to false in test pipeline

Author: billy-moore-98

.github/workflows/test_and_deploy.yml

@@ -35,7 +35,7 @@ jobs:
 
     - name: Check code formatting with Ruff
       run: ruff format --diff --target-version=py39
-      continue-on-error: true
+      continue-on-error: false
 
     - name: Run tests and generate HTML report
       run: |

dags/fetch.py

@@ -6,70 +6,82 @@
 from airflow.operators.dagrun import TriggerDagRunOperator
 from airflow.operators.empty import EmptyOperator
 from airflow.operators.python import PythonOperator
-from airflow.providers.amazon.aws.operators.lambda_function import AwsLambdaInvokeFunctionOperator
+from airflow.providers.amazon.aws.operators.lambda_function import (
+    AwsLambdaInvokeFunctionOperator,
+)
 
-lambda_function_fetch_name = os.getenv('LAMBDA_FUNCTION_FETCH_NAME')
-lambda_function_validate_name = os.getenv('LAMBDA_FUNCTION_VALIDATE_NAME')
+lambda_function_fetch_name = os.getenv("LAMBDA_FUNCTION_FETCH_NAME")
+lambda_function_validate_name = os.getenv("LAMBDA_FUNCTION_VALIDATE_NAME")
 
 default_args = {
-    'owner': 'Billy Moore',
-    'retries': 1,
-    'retry_delay': datetime.timedelta(minutes=1),
-    'depends_on_past': False,
-    'email_on_failure': False,
-    'email_on_retry': False
+    "owner": "Billy Moore",
+    "retries": 1,
+    "retry_delay": datetime.timedelta(minutes=1),
+    "depends_on_past": False,
+    "email_on_failure": False,
+    "email_on_retry": False,
 }
 
+
 # callable to check the result of lambda functions
 def check_lambda_result(task_id, **context):
-    result = context['ti'].xcom_pull(task_ids=task_id)
+    result = context["ti"].xcom_pull(task_ids=task_id)
     if result is None:
-        raise ValueError(f'Lambda function {task_id} failed to return a result.')
-    payload = result.get('Payload')
+        raise ValueError(f"Lambda function {task_id} failed to return a result.")
+    payload = result.get("Payload")
     if payload:
         response = json.loads(payload.read())
-        if response.get('StatusCode') != 200:
-            raise ValueError(f'Lambda function {task_id} failed with status code: {response.get('StatusCode')}')
+        if response.get("StatusCode") != 200:
+            raise ValueError(
+                f"Lambda function {task_id} failed with status code: {response.get('StatusCode')}"
+            )
     else:
-        raise ValueError(f'Lambda function {task_id} returned no payload.')
+        raise ValueError(f"Lambda function {task_id} returned no payload.")
+
 
 with DAG(
-    dag_id='fetch',
+    dag_id="fetch",
     default_args=default_args,
     catchup=False,
-    schedule_interval='@hourly'
+    schedule_interval="@hourly",
 ) as dag:
-    
-    initiate = EmptyOperator(task_id='initiate')
+    initiate = EmptyOperator(task_id="initiate")
 
     lambda_fetch = AwsLambdaInvokeFunctionOperator(
-        task_id='lambda_fetch',
+        task_id="lambda_fetch",
         function_name=lambda_function_fetch_name,
-        payload={'filename': 'market_data_{{ ts_nodash }}.json'}
+        payload={"filename": "market_data_{{ ts_nodash }}.json"},
     )
 
     check_fetch = PythonOperator(
-        task_id='check_fetch',
+        task_id="check_fetch",
         python_callable=check_lambda_result,
-        op_kwargs={'task_id': 'lambda_fetch'}
+        op_kwargs={"task_id": "lambda_fetch"},
     )
 
     lamda_validate = AwsLambdaInvokeFunctionOperator(
-        task_id='lambda_validate',
+        task_id="lambda_validate",
         function_name=lambda_function_validate_name,
-        payload={'filename': 'market_data_{{ ts_nodash }}.json'}
+        payload={"filename": "market_data_{{ ts_nodash }}.json"},
     )
 
     check_validate = PythonOperator(
-        task_id='check_validate',
+        task_id="check_validate",
         python_callable=check_lambda_result,
-        op_kwargs={'task_id': 'lambda_validate'}
+        op_kwargs={"task_id": "lambda_validate"},
     )
 
     trigger_snowflake_ingestion = TriggerDagRunOperator(
-        task_id='trigger_snowflake_ingestion',
-        trigger_dag_id='ingest',
-        conf={'execution_timestamp': '{{ ts_nodash }}'},
+        task_id="trigger_snowflake_ingestion",
+        trigger_dag_id="ingest",
+        conf={"execution_timestamp": "{{ ts_nodash }}"},
     )
 
-    initiate >> lambda_fetch >> check_fetch >> lamda_validate >> check_validate >> trigger_snowflake_ingestion
+    (
+        initiate
+        >> lambda_fetch
+        >> check_fetch
+        >> lamda_validate
+        >> check_validate
+        >> trigger_snowflake_ingestion
+    )

dags/ingest.py

@@ -5,46 +5,62 @@
 from pathlib import Path
 
 default_args = {
-    'owner': 'Billy Moore',
-    'retries': 1,
-    'retry_delay': datetime.timedelta(minutes=1),
-    'depends_on_past': False,
-    'email_on_failure': False,
-    'email_on_retry': False
+    "owner": "Billy Moore",
+    "retries": 1,
+    "retry_delay": datetime.timedelta(minutes=1),
+    "depends_on_past": False,
+    "email_on_failure": False,
+    "email_on_retry": False,
 }
 
+
 def read_sql_query(dir: str, name: str) -> str:
     root_dir = Path(__file__).parent.parent
-    sql_path = root_dir / 'sql' / dir / name
-    with open(sql_path, 'r') as f:
+    sql_path = root_dir / "sql" / dir / name
+    with open(sql_path, "r") as f:
         sql = f.read()
     return sql
 
-with DAG(
-    dag_id='ingest',
-    default_args=default_args
-):
+
+with DAG(dag_id="ingest", default_args=default_args):
     execution_timestamp = "{{ dag_run.conf['execution_timestamp'] }}"
 
-    def snowflake_task_factory(task_id: str, filename: str, timestamp_param: bool = False):
+    def snowflake_task_factory(
+        task_id: str, filename: str, timestamp_param: bool = False
+    ):
         return SnowflakeOperator(
             task_id=task_id,
-            sql=read_sql_query('loading', filename),
-            snowflake_conn_id='snowflake_predictit',
-            params={
-                'execution_timestamp': execution_timestamp
-            } if timestamp_param else {}
+            sql=read_sql_query("loading", filename),
+            snowflake_conn_id="snowflake_predictit",
+            params={"execution_timestamp": execution_timestamp}
+            if timestamp_param
+            else {},
         )
 
-    load_stage_raw = snowflake_task_factory('load_stage_raw', 'load_stage_raw.sql', timestamp_param=True)
-    load_stg_dim_markets = snowflake_task_factory('load_stg_dim_markets', 'load_stg_dim_markets.sql')
-    load_dim_markets = snowflake_task_factory('load_dim_markets', 'load_dim_markets.sql')
-    load_stg_dim_contracts = snowflake_task_factory('load_stg_dim_contracts', 'load_stg_dim_contracts.sql')
-    load_dim_contracts = snowflake_task_factory('load_dim_contracts', 'load_dim_contracts.sql')
-    load_fact_prices = snowflake_task_factory('load_fact_prices', 'load_fact_prices.sql')
-
-    load_stage_raw >> load_stg_dim_markets >> load_dim_markets >> load_stg_dim_contracts >> load_dim_contracts >> load_fact_prices
-
-    
+    load_stage_raw = snowflake_task_factory(
+        "load_stage_raw", "load_stage_raw.sql", timestamp_param=True
+    )
+    load_stg_dim_markets = snowflake_task_factory(
+        "load_stg_dim_markets", "load_stg_dim_markets.sql"
+    )
+    load_dim_markets = snowflake_task_factory(
+        "load_dim_markets", "load_dim_markets.sql"
+    )
+    load_stg_dim_contracts = snowflake_task_factory(
+        "load_stg_dim_contracts", "load_stg_dim_contracts.sql"
+    )
+    load_dim_contracts = snowflake_task_factory(
+        "load_dim_contracts", "load_dim_contracts.sql"
+    )
+    load_fact_prices = snowflake_task_factory(
+        "load_fact_prices", "load_fact_prices.sql"
+    )
 
-    
+    (
+        load_stage_raw
+        >> load_stg_dim_markets
+        >> load_dim_markets
+        >> load_stg_dim_contracts
+        >> load_dim_contracts
+        >> load_fact_prices
+    )

lambda_fetch/lambda_function.py

@@ -9,7 +9,8 @@
 logger.setLevel(logging.INFO)
 
 predictit = PredictitAPI()
-s3_client = boto3.client('s3')
+s3_client = boto3.client("s3")
+
 
 def lambda_function(filename: str) -> None:
     """
@@ -18,15 +19,15 @@ def lambda_function(filename: str) -> None:
         excution_timestamp (str): The execution timestamp for the data
     """
     # Poll the PredictIt API
-    logger.info('Polling PredictIt API market data now')
+    logger.info("Polling PredictIt API market data now")
     data = predictit.poll_market_data()
-    logger.info('Successfully polled API')
-    logger.info('Storing to S3 now')
-    bucket = os.getenv('S3_BUCKET')
+    logger.info("Successfully polled API")
+    logger.info("Storing to S3 now")
+    bucket = os.getenv("S3_BUCKET")
     if not bucket:
         raise ValueError("S3_BUCKET environment variable is not set")
     predictit.store_to_s3(data, bucket=bucket, filename=filename)
-    logging.info('Successfully stored data to S3')
+    logging.info("Successfully stored data to S3")
 
 
 def lambda_handler(event, context) -> Optional[dict]:
@@ -40,14 +41,14 @@ def lambda_handler(event, context) -> Optional[dict]:
         Dict status message
     """
     try:
-        filename = event.get('filename')
+        filename = event.get("filename")
         if not filename:
             raise ValueError("Filename must be provided in the event data")
         lambda_function(filename)
         return {
-            'StatusCode': 200,
-            'message': 'PredictAPI data succcessfully polled and stored to S3'
+            "StatusCode": 200,
+            "message": "PredictAPI data succcessfully polled and stored to S3",
         }
     except Exception as e:
-        logger.error(f'Error occurred: {e}')
-        raise
+        logger.error(f"Error occurred: {e}")
+        raise

lambda_validate/lambda_function.py

@@ -9,7 +9,8 @@
 logger = logging.getLogger()
 logger.setLevel(logging.INFO)
 
-s3_client = boto3.client('s3')
+s3_client = boto3.client("s3")
+
 
 def lambda_function(execution_timestamp: str):
     """
@@ -19,32 +20,32 @@ def lambda_function(execution_timestamp: str):
         execution_timestamp (str): The execution timestamp for the data
     """
     # Validate the PredictIt API data
-    logger.info('Validating PredictIt API data now')
-    bucket = os.getenv('S3_BUCKET')
+    logger.info("Validating PredictIt API data now")
+    bucket = os.getenv("S3_BUCKET")
     if not bucket:
         raise ValueError("S3_BUCKET environment variable is not set")
-    source_key = f'predictit/stage/market_data_{execution_timestamp}.json'
-    destination_key = f'predictit/raw_data/market_data_{execution_timestamp}.json'
+    source_key = f"predictit/stage/market_data_{execution_timestamp}.json"
+    destination_key = f"predictit/raw_data/market_data_{execution_timestamp}.json"
     # Load the data from S3
     s3_object = s3_client.get_object(Bucket=bucket, Key=source_key)
-    data = json.loads(s3_object['Body'].read())
-    
+    data = json.loads(s3_object["Body"].read())
+
     try:
         # Validate the data
         PredictitResponse(**data)
-        logger.info('Successfully validated data')
+        logger.info("Successfully validated data")
     except Exception as e:
-        logger.error(f'Error occurred during data validation: {e}')
+        logger.error(f"Error occurred during data validation: {e}")
         raise
 
     # copy to raw data and delete stage data
     s3_client.copy_object(
         Bucket=bucket,
-        CopySource={'Bucket': bucket, 'Key': source_key},
-        Key=destination_key
+        CopySource={"Bucket": bucket, "Key": source_key},
+        Key=destination_key,
     )
     s3_client.delete_object(Bucket=bucket, Key=source_key)
-        
+
 
 def lambda_handler(event, context) -> Optional[dict]:
     """
@@ -57,14 +58,11 @@ def lambda_handler(event, context) -> Optional[dict]:
         Dict status message
     """
     try:
-        execution_timestamp = event.get('execution_timestamp')
+        execution_timestamp = event.get("execution_timestamp")
         if not execution_timestamp:
             raise ValueError("Execution timestamp must be provided in the event data")
         lambda_function(execution_timestamp)
-        return {
-            'StatusCode': 200,
-            'message': 'PredictAPI data successfully validated'
-        }
+        return {"StatusCode": 200, "message": "PredictAPI data successfully validated"}
     except Exception as e:
-        logger.error(f'Error occurred: {e}')
+        logger.error(f"Error occurred: {e}")
         raise

src/api.py

@@ -10,11 +10,11 @@
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
-class PredictitAPI:
 
+class PredictitAPI:
     def __init__(self, base_url="https://www.predictit.org/api/marketdata"):
         self.base_url = base_url
-        self.s3_client = boto3.client('s3')
+        self.s3_client = boto3.client("s3")
 
     def poll_market_data(self, market_id: Optional[str] = None) -> Optional[dict]:
         """
@@ -41,19 +41,21 @@ def poll_market_data(self, market_id: Optional[str] = None) -> Optional[dict]:
         except Exception as e:
             logging.error(f"An error occurred: {e}")
 
-    def store_to_s3(self, data: dict, bucket: Optional[str] = None, filename: Optional[str] = None):
+    def store_to_s3(
+        self, data: dict, bucket: Optional[str] = None, filename: Optional[str] = None
+    ):
         if not filename:
-            timestamp = datetime.datetime.utcnow().strftime('%Y-%m-%dT%H-%M-%S')
-            filename = f'market_data_{timestamp}.json'
-        key = f'predictit/stage/{filename}'
+            timestamp = datetime.datetime.utcnow().strftime("%Y-%m-%dT%H-%M-%S")
+            filename = f"market_data_{timestamp}.json"
+        key = f"predictit/stage/{filename}"
         try:
             self.s3_client.put_object(
                 Bucket=bucket,
                 Key=key,
                 Body=json.dumps(data),
-                ContentType='application/json'
+                ContentType="application/json",
             )
-            logging.info(f'Uploaded {key} to S3 bucket {bucket}')
+            logging.info(f"Uploaded {key} to S3 bucket {bucket}")
         except ClientError as e:
-            logging.error(f'Failed to upload to bucket: {e}')
-            raise
+            logging.error(f"Failed to upload to bucket: {e}")
+            raise