def test_env_var_priority(self): c = SqliteHook.get_connection(conn_id='airflow_db') self.assertNotEqual('ec2.compute.com', c.host) with mock.patch.dict('os.environ', { 'AIRFLOW_CONN_AIRFLOW_DB': 'postgres://*****:*****@ec2.compute.com:5432/the_database', }): c = SqliteHook.get_connection(conn_id='airflow_db') self.assertEqual('ec2.compute.com', c.host) self.assertEqual('the_database', c.schema) self.assertEqual('username', c.login) self.assertEqual('password', c.password) self.assertEqual(5432, c.port)
def test_using_env_var(self): c = SqliteHook.get_connection(conn_id='test_uri') self.assertEqual('ec2.compute.com', c.host) self.assertEqual('the_database', c.schema) self.assertEqual('username', c.login) self.assertEqual('password', c.password) self.assertEqual(5432, c.port)
def test_using_unix_socket_env_var(self): c = SqliteHook.get_connection(conn_id='test_uri_no_creds') self.assertEqual('ec2.compute.com', c.host) self.assertEqual('the_database', c.schema) self.assertIsNone(c.login) self.assertIsNone(c.password) self.assertIsNone(c.port)
def store_recording(**context): """Stores a random integer from 1 to 10 with timestamp into data table""" conn = SqliteHook.get_connection( context['params']['conn_id']).get_hook().get_conn() time = context['execution_date'] values = (time, random.randint(1, 10)) logging.info('Inserting: {0:%Y-%m-%d %H:%M:%S}, {1}'.format(*values)) with conn: cur = conn.cursor() cur.execute("INSERT INTO data VALUES (?, ?)", values)
def calc_hourly_totals(**context): """Calculates sum of values from data table for the previous hour and puts it into totals table""" conn = SqliteHook.get_connection( context['params']['conn_id']).get_hook().get_conn() hour = context['execution_date'] prev_hour = hour - timedelta(hours=1) with conn: cur = conn.cursor() cur.execute( "SELECT count(*) FROM data " "WHERE strftime('%Y-%m-%d %H', time) = strftime('%Y-%m-%d %H', ?)", (prev_hour, )) data_count = cur.fetchone()[0] logging.info('Found {} values for {}'.format(data_count, hour)) if data_count != 4: logging.warn('Need all 4 values to calculate hourly value.') raise AirflowException( "Need all 4 values to calculate hourly value.") cur.execute( "SELECT count(*) FROM hourly " "WHERE hour = strftime('%Y-%m-%d %H:00:00', ?)", (hour, )) hour_count = cur.fetchone()[0] if hour_count: logging.warn( 'There is an hourly value for {} already.'.format(hour)) return cur.execute( "INSERT INTO hourly SELECT strftime('%Y-%m-%d %H:00:00', ?), sum(value) FROM data " "WHERE strftime('%Y-%m-%d %H', time) = strftime('%Y-%m-%d %H', ?)", (hour, prev_hour)) logging.info('Inserted hourly value for {}.'.format(hour))
def create_tables(**context): hook = SqliteHook.get_connection(context['params']['conn_id']).get_hook() run_sql(hook, "CREATE TABLE IF NOT EXISTS data (time TIMESTAMP, value INT);") run_sql(hook, "CREATE TABLE IF NOT EXISTS hourly (hour TIMESTAMP, total INT);")