incubator-airflow/airflow/hooks/postgres_hook.py

73 строки
2.5 KiB
Python

# -*- 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}``.
Note: For Redshift, use keepalives_idle in the extra connection parameters
and set it to less than 300 seconds.
"""
conn_name_attr = 'postgres_conn_id'
default_conn_name = 'postgres_default'
supports_autocommit = True
def __init__(self, *args, **kwargs):
super(PostgresHook, self).__init__(*args, **kwargs)
self.schema = kwargs.pop("schema", None)
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=self.schema or 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', 'application_name',
'keepalives_idle']:
conn_args[arg_name] = arg_val
psycopg2_conn = psycopg2.connect(**conn_args)
return psycopg2_conn
@staticmethod
def _serialize_cell(cell, conn):
"""
Postgresql will adapt all arguments to the execute() method internally,
hence we return cell without any conversion.
See http://initd.org/psycopg/docs/advanced.html#adapting-new-types for
more information.
:param cell: The cell to insert into the table
:type cell: object
:param conn: The database connection
:type conn: connection object
:return: The cell
:rtype: object
"""
return cell