Coverage for airflow/hooks/postgres_hook.py : 47%

# -*- 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 psycopg2 

import psycopg2.extensions 

from airflow.hooks.dbapi_hook import DbApiHook 

class PostgresHook(DbApiHook): 

''' 

Interact with Postgres. 

You can specify ssl parameters in the extra field of your connection 

as ``{"sslmode": "require", "sslcert": "/path/to/cert.pem", etc}``. 

''' 

conn_name_attr = 'postgres_conn_id' 

default_conn_name = 'postgres_default' 

supports_autocommit = False 

def get_conn(self): 

conn = self.get_connection(self.postgres_conn_id) 

conn_args = dict( 

host=conn.host, 

user=conn.login, 

password=conn.password, 

dbname=conn.schema, 

port=conn.port) 

# check for ssl parameters in conn.extra 

for arg_name, arg_val in conn.extra_dejson.items(): 

if arg_name in ['sslmode', 'sslcert', 'sslkey', 'sslrootcert', 'sslcrl']: 

conn_args[arg_name] = arg_val 

psycopg2_conn = psycopg2.connect(**conn_args) 

if psycopg2_conn.server_version < 70400: 

self.supports_autocommit = True 

return psycopg2_conn 

@staticmethod 

def _serialize_cell(cell): 

return psycopg2.extensions.adapt(cell).getquoted().decode('utf-8')