from airflow.models.baseoperator import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.hooks.S3_hook import S3Hook

import logging
class S3SearchFilingsOperator(BaseOperator):
	"""
	Queries the Datastore API and uploads the processed info as a csv to the S3 bucket. 

	:param source_s3_bucket:          Choose source s3 bucket
	:param source_s3_directory:       Source s3 directory
	:param s3_conn_id:	              S3 Connection ID
	:param destination_s3_bucket:		  S3 Bucket Destination
	"""

	@apply_defaults
	def __init__(
		self,
		source_s3_bucket=None,
		source_s3_directory=True,
		s3_conn_id=True,
		destination_s3_bucket=None,
		destination_s3_directory=None,
		search_terms=[],
		*args, 
		**kwargs) -> None: 
		
		super().__init__(*args, **kwargs)
		
		self.source_s3_bucket = source_s3_bucket
		self.source_s3_directory = source_s3_directory
		self.s3_conn_id = s3_conn_id
		self.destination_s3_bucket = destination_s3_bucket	
		self.destination_s3_directory = destination_s3_directory	

	def execute(self, context):
		"""
		Executes the operator.
		"""
		s3_hook = S3Hook(self.s3_conn_id)
		keys = s3_hook.list_keys(bucket_name=self.source_s3_bucket)
		for key in keys:
			# download file
			filename=s3_hook.download_file(bucket_name=self.source_s3_bucket, key=key, local_path="./data")
			logging.info(filename)
			with open(filename, 'rb') as handle:
				filing = handle.read()
			filing = pickle.loads(filing)
			logging.info(filing.keys())

from keywordSearch.operators.s3_search_filings_operator import S3SearchFilingsOperator

from airflow import DAG
from airflow.utils.dates import days_ago

from datetime import timedelta

# from aws_pull import aws_pull

default_args = {
  "owner" : "airflow",
  "depends_on_past" : False,
  "start_date": days_ago(2),
  "email" : ["[email protected]"],
  "email_on_failure" : False,
  "email_on_retry" : False,
  "retries" : 1,
  "retry_delay": timedelta(seconds=30)
}

with DAG("keyword-search-full-load",
  default_args=default_args,
  description="Syntax Keyword Search",
  max_active_runs=1,
  schedule_interval=None) as dag:

  op3 = S3SearchFilingsOperator(
    task_id="s3_search_filings",
    source_s3_bucket="processed-filings",
    source_s3_directory="citations",
    s3_conn_id="Syntax_S3",
    destination_s3_bucket="keywordsearch",
    destination_s3_directory="results",
    dag=dag
  )

op3