Coverage for airflow/operators/postgres_operator.py : 58%

# -*- coding: utf-8 -*- 

# 

# Licensed under the Apache License, Version 2.0 (the "License"); 

# you may not use this file except in compliance with the License. 

# You may obtain a copy of the License at 

# 

# http://www.apache.org/licenses/LICENSE-2.0 

# 

# Unless required by applicable law or agreed to in writing, software 

# distributed under the License is distributed on an "AS IS" BASIS, 

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

# See the License for the specific language governing permissions and 

# limitations under the License. 

import logging 

from airflow.hooks.postgres_hook import PostgresHook 

from airflow.models import BaseOperator 

from airflow.utils.decorators import apply_defaults 

class PostgresOperator(BaseOperator): 

""" 

Executes sql code in a specific Postgres database 

:param postgres_conn_id: reference to a specific postgres database 

:type postgres_conn_id: string 

:param sql: the sql code to be executed 

:type sql: Can receive a str representing a sql statement, 

a list of str (sql statements), or reference to a template file. 

Template reference are recognized by str ending in '.sql' 

""" 

template_fields = ('sql',) 

template_ext = ('.sql',) 

ui_color = '#ededed' 

@apply_defaults 

def __init__( 

self, sql, 

postgres_conn_id='postgres_default', autocommit=False, 

parameters=None, 

*args, **kwargs): 

super(PostgresOperator, self).__init__(*args, **kwargs) 

self.sql = sql 

self.postgres_conn_id = postgres_conn_id 

self.autocommit = autocommit 

self.parameters = parameters 

def execute(self, context): 

logging.info('Executing: ' + str(self.sql)) 

self.hook = PostgresHook(postgres_conn_id=self.postgres_conn_id) 

self.hook.run(self.sql, self.autocommit, parameters=self.parameters)