Source code for nsaph.db

"""
API to establish database connection

Connection details and credentials are specified in
database.ini file

This module supports connecting via ssh tunnel.
This happens automatically if the given section of
database.ini contains ssh_user key.

"""
import json
#  Copyright (c) 2021. Harvard University
#
#  Developed by Research Software Engineering,
#  Faculty of Arts and Sciences, Research Computing (FAS RC)
#  Author: Michael A Bouzinier
#
#  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
import socket
import paramiko
import psycopg2
import os
import sshtunnel
from configparser import ConfigParser
import boto3
from botocore.exceptions import ClientError

from deprecated.sphinx import deprecated

from nsaph import app_name


[docs]class Connection: default_filename = 'database.ini' default_section = 'postgresql' aws_default_region = "us-east-1" aws_default_secret_name = "nsaph/public/dorieh/"
[docs] @classmethod def read_config(cls, filename, section): home = os.getenv("NSAPH_HOME") if home and not os.path.isabs(filename): filename = os.path.join(home, filename) if not os.path.isfile(filename): raise ValueError( "File {} does not exist or is not readable".format(filename) ) parser = ConfigParser() parser.read(filename) parameters = {} if parser.has_section(section): params = parser.items(section) for param in params: parameters[param[0]] = param[1] else: raise ValueError('Section {0} not found in the {1} file' .format(section, filename)) return parameters
[docs] @classmethod def get_aws_secret(cls, region_name = None, secret_name = None): if secret_name is None: secret_name = cls.aws_default_secret_name if region_name is None: region_name = cls.aws_default_region # Create a Secrets Manager client session = boto3.session.Session() client = session.client( service_name='secretsmanager', region_name=region_name ) try: get_secret_value_response = client.get_secret_value( SecretId=secret_name ) except ClientError as e: # For a list of exceptions thrown, see # https://docs.aws.amazon.com/secretsmanager/latest/apireference/API_GetSecretValue.html raise e # Decrypts secret using the associated KMS key. secret = get_secret_value_response['SecretString'] return secret
# Your code goes here.
[docs] @staticmethod def host_name(): return socket.gethostname().lower()
[docs] @classmethod def resolve_host(cls,host): hosts = host.lower().split(':') if (len(hosts) < 2): return host hostname = cls.host_name().lower() if (hostname in hosts[1:]): return "localhost" return hosts[0]
[docs] @staticmethod def default_port() -> int: return 5432
def __init__(self, filename=None, section=None, silent: bool = False, app_name_postfix = ""): if not filename: filename = Connection.default_filename if not section: section = Connection.default_section self.parameters = self.read_config(filename, section) if "secret" in self.parameters and self.parameters["secret"].startswith("aws:"): pp = self.parameters["secret"].split(':') region = self.aws_default_region name = self.aws_default_secret_name for x in pp: xx = x.split('=') if xx[0] == "region": region = xx[1] elif xx[0] == "name": name = xx[1] data = json.loads(self.get_aws_secret(region, name)) print(data) del self.parameters["secret"] for key in ["password", "database"]: if key in data: self.parameters[key] = data[key] for key in ["host", "port"]: # to account for port forwarding if key in data and key not in self.parameters: self.parameters[key] = data[key] if "username" in data: self.parameters["user"] = data["username"] if "application_name" not in self.parameters: name = "nsaph:" + app_name() + app_name_postfix self.parameters["application_name"] = name self.connection = None self.tunnel = None self.silent = silent self.types = None
[docs] def pid(self) -> int: return self.get_pid(self.connection)
[docs] @staticmethod def get_pid(connection) -> int: with connection.cursor() as cursor: cursor.execute("SELECT pg_backend_pid()") for row in cursor: return row[0]
[docs] def connect_to_database(self, params): if not self.silent: logging.info('Connecting to the PostgreSQL database...') conn = psycopg2.connect(**params) return conn
[docs] def connect(self, autocommit = None): if "ssh_user" in self.parameters: self.connection = self.connect_via_tunnel() else: self.connection = self.connect_to_database(self.parameters) if autocommit is not None: self.connection.autocommit = autocommit info = self.connection.info pid = self.pid() if not self.silent: logging.info("Connected to: {}@{}:{}/{}[{:d}]" .format(info.user, info.host, info.port, info.dbname, pid)) return self.connection
[docs] def connect_via_tunnel(self): host = self.parameters["host"] home = os.path.expanduser('~') mypkey = paramiko.RSAKey.from_private_key_file(os.path.join(home, ".ssh", "id_rsa")) port = self.parameters.get("port", self.default_port()) self.tunnel = sshtunnel.SSHTunnelForwarder( (host, 22), ssh_username=self.parameters["ssh_user"], ssh_pkey=mypkey, remote_bind_address=("localhost", port) ) self.tunnel.start() params = dict() params.update(self.parameters) del params["ssh_user"] params["port"] = self.tunnel.local_bind_port params["host"] = self.tunnel.local_bind_host return self.connect_to_database(params)
[docs] def get_database_types(self): if not self.types: sql = "SELECT oid, typname from pg_catalog.pg_type" cursor = self.connection.cursor() cursor.execute(sql) rows = cursor.fetchall() self.types = { row[0]: row[1] for row in rows } cursor.close() return self.types
[docs] def close(self): if (self.connection and not self.connection.closed): self.connection.close() self.connection = None if (self.tunnel): self.tunnel.stop() self.tunnel = None
def __exit__(self, exc_type, exc_val, exc_tb): self.close() def __enter__(self): return self.connect()
[docs]def test_connection (): with Connection() as conn: cur = conn.cursor() logging.info('PostgreSQL database version:') cur.execute('SELECT version()') db_version = cur.fetchone() logging.info(db_version) cur.close() logging.info('Database connection closed.')
[docs]@deprecated( reason="Use psycopg2.extras.RealDictCursor", version="0.2" ) class ResultSetDeprecated: SIZE = 10000 def __init__(self, cursor, metadata: dict): self.cursor = cursor description = self.cursor.description self.header = [c.name for c in description] self.types = [metadata[c.type_code] for c in description] self.rows = self.cursor.fetchmany(self.SIZE) self.idx = 0 def __iter__(self): return self def __next__(self): self.idx += 1 if self.idx > len(self.rows): self.rows = self.cursor.fetchmany(self.SIZE) if len(self.rows) < 1: raise StopIteration self.idx = 1 row = self.rows[self.idx - 1] return {self.header[i]: row[i] for i in range(len(self.header))}
# class PreparedInsert: # def __init__(self, cursor, sql: str): if __name__ == '__main__': test_connection()