def test_parse_dsn(self):
from psycopg2 import ProgrammingError
self.assertEqual(
ext.parse_dsn('dbname=test user=tester password=secret'),
dict(user='******', password='******', dbname='test'),
"simple DSN parsed")
self.assertRaises(ProgrammingError, ext.parse_dsn,
"dbname=test 2 user=tester password=secret")
self.assertEqual(
ext.parse_dsn("dbname='test 2' user=tester password=secret"),
dict(user='******', password='******', dbname='test 2'),
"DSN with quoting parsed")
# Can't really use assertRaisesRegexp() here since we need to
# make sure that secret is *not* exposed in the error messgage
# (and it also requires python >= 2.7).
raised = False
try:
# unterminated quote after dbname:
ext.parse_dsn("dbname='test 2 user=tester password=secret")
except ProgrammingError, e:
raised = True
self.assertTrue(str(e).find('secret') < 0,
"DSN was not exposed in error message")
def test_parse_dsn_uri(self):
self.assertEqual(ext.parse_dsn('postgresql://*****:*****@/test'),
dict(user='******', password='******', dbname='test'),
"valid URI dsn parsed")
raised = False
try:
# extra '=' after port value
ext.parse_dsn(dsn='postgresql://*****:*****@/test?port=1111=x')
except psycopg2.ProgrammingError, e:
raised = True
self.assertTrue(str(e).find('secret') < 0,
"URI was not exposed in error message")
def run_osm2pgsql(options):
""" Run osm2pgsql with the given options.
"""
env = os.environ
cmd = [
options['osm2pgsql'], '--hstore', '--latlon', '--slim',
'--with-forward-dependencies', 'false', '--log-progress', 'true',
'--number-processes',
str(options['threads']), '--cache',
str(options['osm2pgsql_cache']), '--output', 'gazetteer', '--style',
str(options['osm2pgsql_style'])
]
if options['append']:
cmd.append('--append')
if options['flatnode_file']:
cmd.extend(('--flat-nodes', options['flatnode_file']))
dsn = parse_dsn(options['dsn'])
if 'password' in dsn:
env['PGPASSWORD'] = dsn['password']
if 'dbname' in dsn:
cmd.extend(('-d', dsn['dbname']))
if 'user' in dsn:
cmd.extend(('--username', dsn['user']))
for param in ('host', 'port'):
if param in dsn:
cmd.extend(('--' + param, dsn[param]))
cmd.append(str(options['import_file']))
subprocess.run(cmd, cwd=options.get('cwd', '.'), env=env, check=True)
def test_unicode_value(self):
snowman = u"\u2603"
d = ext.parse_dsn('dbname=' + snowman)
if sys.version_info[0] < 3:
self.assertEqual(d['dbname'], snowman.encode('utf8'))
else:
self.assertEqual(d['dbname'], snowman)
def before_call(self, scope, cursor, connection, _args, _kwargs, response, exception):
span = scope.span
span.domain_name = constants.DomainNames['DB']
span.class_name = constants.ClassNames['POSTGRESQL']
dsn = parse_dsn(connection.dsn)
query = ''
operation = ''
try:
query = _args[0]
if len(query) > 0:
operation = query.split()[0].strip("\"").lower()
except Exception:
pass
tags = {
constants.SpanTags['OPERATION_TYPE']: PostgreIntegration._OPERATION_TO_TYPE.get(operation, ''),
constants.SpanTags['DB_INSTANCE']: dsn.get('dbname', ''),
constants.SpanTags['DB_HOST']: dsn.get('host', ''),
constants.SpanTags['DB_TYPE']: "postgresql",
constants.SpanTags['DB_STATEMENT_TYPE']: operation.upper(),
constants.SpanTags['TRIGGER_CLASS_NAME']: "API",
constants.SpanTags['TOPOLOGY_VERTEX']: True
}
if not ConfigProvider.get(config_names.THUNDRA_TRACE_INTEGRATIONS_RDB_STATEMENT_MASK):
tags[constants.DBTags['DB_STATEMENT']] = query
span.tags = tags
def from_options(
cls,
options: optparse.Values,
dsn: str,
hostname: str,
) -> "Host":
"""Build a host instance from options and dsn
>>> options = optparse.Values(
... defaults = {
... "username": "******",
... "host": "test",
... "port": 5432,
... "dbname": "pg",
... })
>>> Host.from_options(options, "", "test")
Host(hostname='test', user='******', host='test', port=5432, dbname='pg')
>>> Host.from_options(options, "host=/tmp port=5432 user=toto dbname=pgbench", "test")
Host(hostname='test', user='******', host='/tmp', port=5432, dbname='pgbench')
>>> Host.from_options(options, "postgresql://toto@localhost:5432/bench", "test")
Host(hostname='test', user='******', host='localhost', port=5432, dbname='bench')
"""
pdsn = parse_dsn(dsn)
return cls(
hostname,
pdsn.get("user", options.username),
pdsn.get("host", options.host),
int(pdsn.get("port", options.port)),
pdsn.get("dbname", options.dbname),
)
def get_database_password(dsn):
if 'ENCRYPTED_DATABASE_PASSWORD' in os.environ:
# boto returns decrypted as b'bytes' so decode to convert to password string
return decrypt(os.environ['ENCRYPTED_DATABASE_PASSWORD']).decode()
else:
dsn_components = parse_dsn(dsn)
return boto3.client('rds').generate_db_auth_token(dsn_components['host'], 5432, dsn_components['user'])
def wrapper(dsn, connection_factory=None, *args, **kwargs):
config = parse_dsn(dsn)
host = config.get("host")
if host is None:
host = kwargs.get("host")
if host is None:
# Host may be left out to use localhost or
# possibly set using environment variables, nothing
# we can do in either case.
logger.error(
"'host' parameter is not present in call to psycopg2.connect, "
"DNS resolution might not work properly"
)
return fn(dsn, connection_factory, *args, **kwargs)
if re.search(host_regex, host):
logger.debug("Host %s matched SRV regex, resolving", host)
host, port = resolve_srv_record(host, srv_resolver)
config["host"] = host
config["port"] = port
dsn = make_dsn(**config)
return fn(dsn, connection_factory, *args, **kwargs)
def truncate_table(self):
print('Truncating table %s ...' % self.target_table)
conn = psycopg2.connect(**parse_dsn(self.db_config.db_connect_string))
cur = conn.cursor()
cur.execute('TRUNCATE TABLE %s' % self.target_table)
conn.commit()
cur.close()
conn.close()
def test_unicode_value(self):
from psycopg2.extensions import parse_dsn
snowman = u"\u2603"
d = parse_dsn('dbname=' + snowman)
if sys.version_info[0] < 3:
self.assertEqual(d['dbname'], snowman.encode('utf8'))
else:
self.assertEqual(d['dbname'], snowman)
def get_database_settings(options):
"""Get the database setting to use.
It will either read --database-uri from the options, or prompt for it.
When prompting for it, it will default to the maas-test-db URI if the
maas-test-db snap is installed and connected.
"""
database_uri = options.database_uri
test_db_socket = os.path.join(os.environ["SNAP_COMMON"], "test-db-socket")
test_db_uri = f"postgres:///maasdb?host={test_db_socket}&user=maas"
if database_uri is None:
default_uri = None
if os.path.exists(test_db_socket):
default_uri = MAAS_TEST_DB_URI
database_uri = required_prompt(
"Database URI",
default=default_uri,
help_text=ARGUMENTS["database-uri"]["help"],
)
if not database_uri:
database_uri = test_db_uri
# parse_dsn gives very confusing error messages if you pass in
# an invalid URI, so let's make sure the URI is of the form
# postgres://... before calling parse_dsn.
if database_uri != MAAS_TEST_DB_URI and not database_uri.startswith(
"postgres://"):
raise DatabaseSettingsError(
f"Database URI needs to be either '{MAAS_TEST_DB_URI}' or "
"start with 'postgres://'")
if database_uri == MAAS_TEST_DB_URI:
database_uri = test_db_uri
try:
parsed_dsn = parse_dsn(database_uri)
except psycopg2.ProgrammingError as error:
raise DatabaseSettingsError("Error parsing database URI: " +
str(error).strip())
unsupported_params = set(parsed_dsn.keys()).difference(
["user", "password", "host", "dbname", "port"])
if unsupported_params:
raise DatabaseSettingsError(
"Error parsing database URI: Unsupported parameters: " +
", ".join(sorted(unsupported_params)))
if "user" not in parsed_dsn:
raise DatabaseSettingsError(f"No user found in URI: {database_uri}")
if "host" not in parsed_dsn:
parsed_dsn["host"] = "localhost"
if "dbname" not in parsed_dsn:
parsed_dsn["dbname"] = parsed_dsn["user"]
database_settings = {
"database_host": parsed_dsn["host"],
"database_name": parsed_dsn["dbname"],
"database_user": parsed_dsn.get("user", ""),
"database_pass": parsed_dsn.get("password"),
}
if "port" in parsed_dsn:
database_settings["database_port"] = int(parsed_dsn["port"])
return database_settings
def __init__(self, dsn):
parsed = ext.parse_dsn(dsn)
# ca chain downloaded from https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/UsingWithRDS.SSL.html
# or `wget https: //s3.amazonaws.com/rds-downloads/rds-combined-ca-bundle.pem`
self.ssl_root_cert = parsed.get(
"sslrootcert", "./certificates/rds-combined-ca-bundle.pem")
self.logger = logging.getLogger(__name__ + "." +
self.__class__.__name__)
def __init__(self, dsn, region):
parsed = ext.parse_dsn(dsn)
self.host = parsed.get("host", "localhost")
self.port = parsed.get("port", "5432")
self.user = parsed.get("user", "postgres")
self.client = boto3.client('rds', region_name=region)
self.logger = logging.getLogger(__name__ + "." +
self.__class__.__name__)
def connect(db_url: str) -> Any:
"""
Return a database connection based on the input database URL.
@param db_url: The database URL.
@returns: The connection.
"""
db_params = parse_dsn(db_url)
return psycopg2.connect(**db_params)
def __init__(self, connection, cursor, _args, _kwargs, start_time,
exception):
"""
Initialize.
:param connection: The SQL engine the event is using
:param cursor: Cursor object used in the even
:param args: args passed to called function
:param kwargs: kwargs passed to called function
:param start_time: Start timestamp (epoch)
:param exception: Exception (if occurred)
"""
super(DBAPIEvent, self).__init__(start_time)
self.event_id = 'dbapi-{}'.format(str(uuid4()))
# in case of pg instrumentation we extract data from the dsn property
if hasattr(connection, 'dsn'):
dsn = parse_dsn(connection.dsn)
db_name = dsn.get('dbname', '')
host = dsn.get('host', 'local')
query = cursor.query
else:
query = _args[0]
host = connection.extract_hostname
db_name = connection.extract_dbname
self.resource['name'] = db_name if db_name else host
# NOTE: The operation might not be identified properly when
# using 'WITH' clause
splitted_query = query.split()
if not splitted_query:
print_debug('Cannot extract operation from query {}'.format(query))
operation = ''
else:
operation = splitted_query[0].lower()
self.resource['operation'] = operation
# override event type with the specific DB type
self.resource['type'] = database_connection_type(
host, self.RESOURCE_TYPE)
self.resource['metadata'] = {
'Host': host,
'Driver': connection.__class__.__module__.split('.')[0],
'Table Name': self._extract_table_name(query, operation)
}
# for select we always want to save the query
if ((operation == 'select') or (not trace_factory.metadata_only)):
self.resource['metadata']['Query'] = query[:MAX_QUERY_SIZE]
if exception is None:
# Update response data
self.resource['metadata']['Related Rows Count'] = int(
cursor.rowcount)
else:
self.set_exception(exception, traceback.format_exc())
def load(self, features=None, nrows=None):
self.categoricals = {}
self.transforms_log = [[], [], [], []]
import csv
from io import StringIO
path = self.options['data_path']
if isinstance(path, StringIO):
path.seek(0)
self.df = pd.read_csv(path,
encoding='utf-8',
escapechar="\\",
usecols=features,
na_values=['?'],
nrows=nrows)
if self.options.get("targetFeature") in self.df.columns:
self.dropna([self.options["targetFeature"]])
else:
if path.startswith("jdbc:"):
import psycopg2
from psycopg2.extensions import parse_dsn
path = path.replace(
'sslfactory=org.postgresql.ssl.NonValidatingFactory&', '')
ary = path.split('tablename')
path = ary[0]
tablename = ary[1]
dataset_name = tablename
self.dbconn_args = parse_dsn(path[5:])
conn = psycopg2.connect(**self.dbconn_args)
self.df = pd.read_sql("select * from %s" % tablename, con=conn)
else:
path, remote_path = self._check_remote_path()
try:
self.df = self.load_from_file(path,
features=features,
nrows=nrows)
except:
if remote_path:
logging.exception(
"Loading local file failed. Download it again...")
self.options['data_path'] = remote_path
path, remote_path = self._check_remote_path(
force_download=True)
self.df = self.load_from_file(path,
features=features,
nrows=nrows)
else:
raise
self.dataset_name = os.path.basename(path)
if self.options.get("targetFeature") in self.df.columns:
self.dropna([self.options["targetFeature"]])
return self
def parse_connect_string(connect_string):
# According to:
# https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING
try:
parsed_string = parse_dsn(connect_string)
except Exception:
raise Exception("Invalid connect string: " + connect_string)
return parsed_string
def _connection_args_from_dsn(dsn):
params = parse_dsn(dsn)
return [
"-h",
params["host"],
"-U",
params["user"],
"-p",
params["port"],
"--no-password",
]
def _set_attributes_from_cursor(span: trace.Span, vendor, cursor):
"""Attempt to set db connection attributes by introspecting the cursor."""
if vendor == "postgres":
# pylint: disable=import-outside-toplevel
from psycopg2.extensions import parse_dsn
if hasattr(cursor, "connection") and hasattr(cursor.connection, "dsn"):
dsn = getattr(cursor.connection, "dsn", None)
if dsn:
data = parse_dsn(dsn)
span.set_attribute(_DB, data.get("dbname"))
span.set_attribute(_HOST, data.get("host"))
span.set_attribute(_PORT, int(data.get("port")))
def args(self, **override):
if self.url:
args = parse_dsn(self.url)
else:
args = {}
for k in ['host', 'port', 'dbname', 'user', 'password']:
v = getattr(self, k)
if v is not None:
args[k] = v
if not args['dbname']:
args['dbname'] = 'tusker'
args.update(override)
return args
def _get_attributes_from_cursor(vendor, cursor, attrs):
"""Attempt to set db connection attributes by introspecting the cursor."""
if vendor == "postgresql":
# pylint: disable=import-outside-toplevel
from psycopg2.extensions import parse_dsn
if hasattr(cursor, "connection") and hasattr(cursor.connection, "dsn"):
dsn = getattr(cursor.connection, "dsn", None)
if dsn:
data = parse_dsn(dsn)
attrs[SpanAttributes.DB_NAME] = data.get("dbname")
attrs[SpanAttributes.NET_PEER_NAME] = data.get("host")
attrs[SpanAttributes.NET_PEER_PORT] = int(data.get("port"))
return attrs
def __init__(self,
source_dsn,
slot_name=None,
keepslot=False,
decoder='wal2json'):
self.source_dsn = source_dsn
self.dbname = parse_dsn(source_dsn)['dbname']
self.conn = None
self.slot_name = slot_name or 'connemara_%s' % self.dbname
self.keepslot = keepslot
self.decoder = decoder
self.statements = []
self.conn = psycopg2.connect(
self.source_dsn, connection_factory=LogicalReplicationConnection)
def purge_old(self):
logging.debug('purge_old()')
# purge all records flagged as upload_record = false
connection = psycopg2.connect(
**parse_dsn(self.db_config.db_connect_string))
cursor = connection.cursor()
# delete all rows
cursor.execute("ROLLBACK;")
cursor.execute("DELETE FROM %s WHERE upload_record != true;" %
self.target_table)
# reset sequence to 0
# cursor.execute("ROLLBACK;")
# seq = self.target_table + '_id_seq'
# cursor.execute("ALTER SEQUENCE %s RESTART;" % seq)
# persist changes
connection.commit()
cursor.close()
connection.close()
def import_events(event, __):
logger = logging.getLogger('event-recorder')
logger.setLevel(logging.INFO)
dsn = os.environ['DB_CONNECTION_STRING']
database_password = None
if 'ENCRYPTED_DATABASE_PASSWORD' in os.environ:
# boto returns decrypted as b'bytes' so decode to convert to password string
database_password = decrypt(
os.environ['ENCRYPTED_DATABASE_PASSWORD']).decode()
else:
dsn_components = parse_dsn(dsn)
database_password = boto3.client('rds').generate_db_auth_token(
dsn_components['host'], 5432, dsn_components['user'])
db_connection = create_db_connection(dsn, database_password)
logger.info('Created connection to DB')
for record in event['Records']:
bucket = record['s3']['bucket']['name']
filename = record['s3']['object']['key']
iterable = fetch_import_file(bucket, filename)
for line in iterable:
try:
message_envelope = json.loads(line)
event = event_from_json_object(message_envelope['document'])
if write_audit_event_to_database(event, db_connection):
if event.event_type == 'session_event' and event.details.get(
'session_event_type') == 'idp_authn_succeeded':
write_billing_event_to_database(event, db_connection)
if event.event_type == 'session_event' and event.details.get(
'session_event_type') == 'fraud_detected':
write_fraud_event_to_database(event, db_connection)
except Exception as exception:
logger.exception('Failed to store message{}'.format(exception))
delete_import_file(bucket, filename)
def run(self, conn, conninfo):
if not conninfo['standby']:
return []
try:
# Get primary parameters from primary_conninfo
dsn = parse_dsn(get_primary_conninfo(conn))
p_host = dsn['host']
# pg_stat_wal_receiver lookup
rows = conn.query("""\
SELECT '{p_host}' AS upstream,
CASE WHEN COUNT(*) > 0 THEN 1 ELSE 0 END AS connected
FROM pg_stat_wal_receiver
WHERE status='streaming' AND
conninfo LIKE '%host={p_host}%'
""".format(p_host=p_host))
return list(rows)
except Exception as e:
logger.exception(str(e))
return []
def connect(self, schema=None):
"""
Method which attempts connection to postresql instance using either
a series of dsn kwargs or a libpq string. If schema is passed it
will set path to that schema so all subsequent queries will be
executed there.
- *schema*: schema to connect to. If not provided, then public.
It returns a Psycpog2 connection class. If it cannot connect it will
return None.
See here for conn class https://www.psycopg.org/docs/connection.html
"""
try:
if self.libpq_string is not None:
connection = psycopg2.connect(**parse_dsn(self.libpq_string))
else:
connection = psycopg2.connect(**self.kwargs)
except DatabaseError as error:
raise DatabaseError(psycopg2_exception_enhanced(error))
if connection is not None and schema is not None:
with connection:
with connection.cursor() as cur:
schema = sql.Identifier(schema)
try:
cur.execute(
sql.SQL(
"create schema if not exists {schema}; "
"set search_path to {schema}, public;").format(
schema=schema))
except Exception as error:
raise DatabaseError(psycopg2_exception_enhanced(error))
return connection
def __init__(self, dsn, region):
self.logger = logging.getLogger(__name__ + "." +
self.__class__.__name__)
parsed = ext.parse_dsn(dsn)
user = parsed.get("user")
password = parsed.get("password")
if not password and user and user.startswith("iam_"):
self.logger.info(
"Using RDS dns kwargs supplier because password is empty and username starts with iam_"
)
self.kwargs = {
'password': RdsIamAuthTokenGenerator(dsn, region),
'sslmode': RdsIamAuthSslModeProvider(dsn),
'sslrootcert': RdsIamAuthSslRootProvider(dsn)
}
else:
self.logger.info(
"Not using RDS dns kwargs supplier because password is not empty or "
"username does not start with iam_")
self.kwargs = {}
def collect_psycopg2_metrics(context, trace, instance, args):
try:
from psycopg2.extensions import parse_dsn
except ImportError: # pragma: no cover
from .dbapi import parse_dsn
connection = instance.connection_proxy
dsn = parse_dsn(connection.dsn)
db = dsn.get("dbname")
hostname = dsn.get("host", "localhost")
port = dsn.get("port", 5432)
command, table = None, None
query = sqlparse.parse(instance.query)
if not query and args:
query = sqlparse.parse(args[0])
if query:
query = query[0]
command = query.get_type()
table = query.get_name()
request = Request(
command=ensure_utf8(command),
key=None,
hostname=ensure_utf8(hostname),
port=ensure_utf8(port),
connectionName=None,
db=ensure_utf8(db),
table=ensure_utf8(table),
)
request = request._asdict()
context.iopipe.mark.db_trace(trace, "postgresql", request)
def get_conn_id(conn):
"""
Extract application_name from dsn
"""
parsed = parse_dsn(conn.dsn)
return parsed['application_name']
# Build the list of objects to ignore.
# This include every object belonging to an extension,
# plus any object specifically ignored in the config file.
objects_in_extensions = dumper.fetch_objects_in_extensions()
all_objects = list(itertools.chain(*objects_in_extensions.values()))
for obj in conf.get('ignored_objects', []):
all_objects.append(obj)
ignore_whole_schemas = conf.get('ignored_schemas', [])
ignore_whole_schemas.append('_timescaledb_internal')
dumper.create_repl_slot()
logger.debug("dumper start")
ddlscript = dumper.dump(ignored_schemas=ignore_whole_schemas)
logger.debug("dumper end")
dbname = parse_dsn(source_dsn)['dbname']
schema_map = {'public': '%s_public' % dbname}
for schema in ddlscript.objects_created_by_type[ObjectType.OBJECT_SCHEMA]:
schema_map[schema] = '%s_%s' % (dbname, schema)
table_mapping = {}
# Remap the object and keep track of tables
# This is fine for «small» dumps, but if we ever get gigantic schemas
# RAM could be an issue.
# In that case we shouldn't keep the whole dump's parse tree in RAM.
filtered_statements = []
for statement in ddlscript.statements:
obj_creation = object_creation(statement)
if obj_creation:
objtype, fqname = obj_creation
if objtype == ObjectType.OBJECT_SCHEMA:
schema_name = fqname
def test_dsn(engine, pg_params):
params = pg_params.copy()
params['password'] = '******'
params['dbname'] = params.pop('database')
params['port'] = str(params['port'])
assert parse_dsn(engine.dsn) == params
