Skip to content

Instantly share code, notes, and snippets.

@ganapathichidambaram
Last active November 11, 2019 09:06
Show Gist options
  • Save ganapathichidambaram/410c099d8138dc05406bfd25290aefa2 to your computer and use it in GitHub Desktop.
Save ganapathichidambaram/410c099d8138dc05406bfd25290aefa2 to your computer and use it in GitHub Desktop.
Default Postgresql Operator Script of Airflow framework
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 sql: the sql code to be executed. (templated)
: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'
:param postgres_conn_id: reference to a specific postgres database
:type postgres_conn_id: str
:param autocommit: if True, each command is automatically committed.
(default value: False)
:type autocommit: bool
:param parameters: (optional) the parameters to render the SQL query with.
:type parameters: mapping or iterable
:param database: name of database which overwrite defined one in connection
:type database: str
"""
template_fields = ('sql',)
template_ext = ('.sql',)
ui_color = '#ededed'
@apply_defaults
def __init__(
self, sql,
postgres_conn_id='postgres_default', autocommit=False,
parameters=None,
database=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
self.database = database
def execute(self, context):
self.log.info('Executing: %s', self.sql)
self.hook = PostgresHook(postgres_conn_id=self.postgres_conn_id,
schema=self.database)
self.hook.run(self.sql, self.autocommit, parameters=self.parameters)
for output in self.hook.conn.notices:
self.log.info(output)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment