Skip to content

Commit 3e866b5

Browse files
author
James William Pye
committed
Add tests for pg.open and identify driver parameters in service files.
This allows search_path to be properly passed in as a setting.
1 parent 43374e1 commit 3e866b5

2 files changed

Lines changed: 114 additions & 3 deletions

File tree

‎postgresql/clientparameters.py‎

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,27 @@ def getuser():
6262
pg_appdata_directory = 'postgresql'
6363
pg_appdata_passfile = 'pgpass.conf'
6464

65+
# In order to support pg_service.conf, it is
66+
# necessary to identify driver parameters, so
67+
# that database configuration parameters can
68+
# be placed in settings.
69+
pg_service_driver_parameters = set([
70+
'user',
71+
'host',
72+
'database',
73+
'port',
74+
'password',
75+
76+
'sslcrtfile',
77+
'sslkeyfile',
78+
'sslrootcrtfile',
79+
'sslrootkeyfile',
80+
81+
'sslmode',
82+
'server_encoding',
83+
'connect_timeout',
84+
])
85+
6586
# environment variables that will be in the parameters' "settings" dictionary.
6687
default_envvar_settings_map = {
6788
'TZ' : 'timezone',
@@ -411,13 +432,20 @@ def x_pg_service(service_name, config):
411432
return
412433

413434
for (k, v) in s:
414-
if k.lower() == 'ldap':
435+
k = k.lower()
436+
if k == 'ldap':
415437
yield ('pg_ldap', ':'.join((k, v)))
416-
elif k.lower() == 'pg_service':
438+
elif k == 'pg_service':
417439
# ignore
418440
pass
419-
elif k.lower() == 'dbname':
441+
elif k == 'hostaddr':
442+
# XXX: should yield ipv as well?
443+
yield (('host',), v)
444+
elif k == 'dbname':
420445
yield (('database',), v)
446+
elif k not in pg_service_driver_parameters:
447+
# it's a GUC.
448+
yield (('settings', k), v)
421449
else:
422450
yield ((k,), v)
423451

‎postgresql/test/test_connect.py‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111

1212
from ..driver import dbapi20 as dbapi20
1313
from .. import driver as pg_driver
14+
from .. import open as pg_open
1415

1516
class test_connect(pg_unittest.TestCaseWithCluster):
1617
"""
@@ -71,6 +72,88 @@ def initialize_database(self):
7172
"""
7273
)
7374

75+
def test_pg_open(self):
76+
# postgresql.open
77+
host, port = self.cluster.address()
78+
# test simple locators..
79+
with pg_open(
80+
'pq://' + 'md5:' + 'md5_password@' + host + ':' + str(port) \
81+
+ '/test'
82+
) as db:
83+
self.failUnlessEqual(db.prepare('select 1')(), [(1,)])
84+
self.failUnless(db.closed)
85+
86+
with pg_open(
87+
'pq://' + 'password:' + 'password_password@' + host + ':' + str(port) \
88+
+ '/test'
89+
) as db:
90+
self.failUnlessEqual(db.prepare('select 1')(), [(1,)])
91+
self.failUnless(db.closed)
92+
93+
with pg_open(
94+
'pq://' + 'trusted@' + host + ':' + str(port) + '/test'
95+
) as db:
96+
self.failUnlessEqual(db.prepare('select 1')(), [(1,)])
97+
self.failUnless(db.closed)
98+
99+
# test environment collection
100+
pgenv = ('PGUSER', 'PGPORT', 'PGHOST', 'PGSERVICE', 'PGPASSWORD', 'PGDATABASE')
101+
stored = list(map(os.environ.get, pgenv))
102+
try:
103+
os.environ.pop('PGSERVICE', None)
104+
os.environ['PGUSER'] = 'md5'
105+
os.environ['PGPASSWORD'] = 'md5_password'
106+
os.environ['PGHOST'] = host
107+
os.environ['PGPORT'] = str(port)
108+
os.environ['PGDATABASE'] = 'test'
109+
# No arguments, the environment provided everything.
110+
with pg_open() as db:
111+
self.failUnlessEqual(db.prepare('select 1')(), [(1,)])
112+
self.failUnlessEqual(db.prepare('select current_user').first(), 'md5')
113+
self.failUnless(db.closed)
114+
finally:
115+
i = 0
116+
for x in stored:
117+
env = pgenv[i]
118+
if x is None:
119+
os.environ.pop(env, None)
120+
else:
121+
os.environ[env] = x
122+
123+
oldservice = os.environ.get('PGSERVICE')
124+
oldsysconfdir = os.environ.get('PGSYSCONFDIR')
125+
try:
126+
with open('pg_service.conf', 'w') as sf:
127+
sf.write('''
128+
[myserv]
129+
user = password
130+
password = password_password
131+
host = {host}
132+
port = {port}
133+
dbname = test
134+
search_path = public
135+
'''.format(host = host, port = port))
136+
sf.flush()
137+
try:
138+
os.environ['PGSERVICE'] = 'myserv'
139+
os.environ['PGSYSCONFDIR'] = os.getcwd()
140+
with pg_open() as db:
141+
self.failUnlessEqual(db.prepare('select 1')(), [(1,)])
142+
self.failUnlessEqual(db.prepare('select current_user').first(), 'password')
143+
self.failUnlessEqual(db.settings['search_path'], 'public')
144+
finally:
145+
if oldservice is None:
146+
os.environ.pop('PGSERVICE', None)
147+
else:
148+
os.environ['PGSERVICE'] = oldservice
149+
if oldsysconfdir is None:
150+
os.environ.pop('PGSYSCONFDIR', None)
151+
else:
152+
os.environ['PGSYSCONFDIR'] = oldsysconfdir
153+
finally:
154+
if os.path.exists('pg_service.conf'):
155+
os.remove('pg_service.conf')
156+
74157
def test_dbapi_connect(self):
75158
host, port = self.cluster.address()
76159
MD5 = dbapi20.connect(

0 commit comments

Comments
 (0)