Skip to content

Commit f293f07

Browse files
author
Alex Malek
committed
Refactor connection code
no args in constructor will use PGHOST if it exists added a 'verbose' flag to see what connections are being tried
1 parent 7878ba3 commit f293f07

File tree

1 file changed

+86
-85
lines changed

1 file changed

+86
-85
lines changed

‎wrds/sql.py

Lines changed: 86 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -30,116 +30,117 @@ class SchemaNotFoundError(FileNotFoundError):
3030

3131

3232
class Connection(object):
33-
def __init__(self, autoconnect=True, **kwargs):
33+
def __init__(self, autoconnect=True, verbose=False, **kwargs):
3434
"""
35-
Set up the connection to the WRDS database.
36-
By default, also establish the connection to the database.
37-
38-
Optionally, the user may specify connection parameters:
39-
*wrds_hostname*: WRDS database hostname
40-
*wrds_port*: database connection port number
41-
*wrds_dbname*: WRDS database name
42-
*wrds_username*: WRDS username
43-
44-
The constructor will use the .pgpass file if it exists.
45-
If not, it will ask the user for a username and password.
46-
It will also direct the user to information on setting up .pgpass.
47-
48-
Additionally, creating the instance will load a list of schemas
49-
the user has permission to access.
50-
51-
:return: None
52-
53-
Usage::
54-
>>> db = wrds.Connection()
55-
Loading library list...
56-
Done
35+
Set up the connection to the WRDS database.
36+
By default, also establish the connection to the database.
37+
38+
Optionally, the user may specify connection parameters:
39+
*wrds_hostname*: WRDS database hostname
40+
*wrds_port*: database connection port number
41+
*wrds_dbname*: WRDS database name
42+
*wrds_username*: WRDS username
43+
*autoconnect*: If false will not immediately establish the connection
44+
45+
The constructor will use the .pgpass file if it exists and may make use of
46+
PostgreSQL environment variables such as PGHOST, PGUSER, etc., if cooresponding
47+
parameters are not set.
48+
If not, it will ask the user for a username and password.
49+
It will also direct the user to information on setting up .pgpass.
50+
51+
Additionally, creating the instance will load a list of schemas
52+
the user has permission to access.
53+
54+
:return: None
55+
56+
Usage::
57+
>>> db = wrds.Connection()
58+
Loading library list...
59+
Done
5760
"""
61+
self._verbose = verbose
5862
self._password = ""
5963
# If user passed in any of these parameters, override defaults.
60-
self._username = kwargs.get("wrds_username", None)
61-
self._hostname = kwargs.get("wrds_hostname", WRDS_POSTGRES_HOST)
64+
self._username = kwargs.get("wrds_username", "")
65+
# PGHOST if set will override default for first attempt
66+
self._hostname = kwargs.get(
67+
"wrds_hostname", os.environ.get('PGHOST', WRDS_POSTGRES_HOST)
68+
)
6269
self._port = kwargs.get("wrds_port", WRDS_POSTGRES_PORT)
6370
self._dbname = kwargs.get("wrds_dbname", WRDS_POSTGRES_DB)
6471
self._connect_args = kwargs.get("wrds_connect_args", WRDS_CONNECT_ARGS)
6572

66-
# If username was passed in, the URI is different.
67-
if self._username:
68-
pguri = "postgresql://{usr}@{host}:{port}/{dbname}"
69-
self.engine = sa.create_engine(
70-
pguri.format(
71-
usr=self._username,
72-
host=self._hostname,
73-
port=self._port,
74-
dbname=self._dbname,
75-
),
76-
isolation_level="AUTOCOMMIT",
77-
connect_args=self._connect_args,
78-
)
79-
# No username passed in, but other parameters might have been.
80-
else:
81-
pguri = "postgresql://{host}:{port}/{dbname}"
82-
self.engine = sa.create_engine(
83-
pguri.format(host=self._hostname, port=self._port, dbname=self._dbname),
84-
isolation_level="AUTOCOMMIT",
85-
connect_args=self._connect_args,
86-
)
8773
if autoconnect:
8874
self.connect()
8975
self.load_library_list()
9076

91-
def connect(self):
92-
"""Make a connection to the WRDS database."""
77+
def __make_sa_engine_conn(self, raise_err = False):
78+
username = self._username
79+
hostname = self._hostname
80+
password = urllib.parse.quote_plus(self._password)
81+
port = self._port
82+
dbname = self._dbname
83+
pguri = f"postgresql://{username}:{password}@{hostname}:{port}/{dbname}"
84+
if self._verbose:
85+
print(f"postgresql://{username}:@{hostname}:{port}/{dbname}")
9386
try:
94-
self.connection = self.engine.connect()
95-
except Exception:
96-
# These things should probably not be exported all over creation
97-
self._username, self._password = self.__get_user_credentials()
98-
pghost = "postgresql://{usr}:{pwd}@{host}:{port}/{dbname}"
9987
self.engine = sa.create_engine(
100-
pghost.format(
101-
usr=self._username,
102-
pwd=urllib.parse.quote_plus(self._password),
103-
host=self._hostname,
104-
port=self._port,
105-
dbname=self._dbname,
106-
),
88+
pguri,
10789
isolation_level="AUTOCOMMIT",
10890
connect_args=self._connect_args,
10991
)
110-
try:
111-
self.connection = self.engine.connect()
112-
except Exception as e:
113-
print("There was an error with your password.")
114-
self._username = None
115-
self._password = None
116-
raise e
117-
118-
# Connection successful. Offer to create a .pgpass for the user.
119-
print("WRDS recommends setting up a .pgpass file.")
120-
do_create_pgpass = ""
121-
while do_create_pgpass != "y" and do_create_pgpass != "n":
122-
do_create_pgpass = input("Create .pgpass file now [y/n]?: ")
123-
124-
if do_create_pgpass == "y":
125-
try:
126-
self.create_pgpass_file()
127-
print("Created .pgpass file successfully.")
128-
except:
129-
print(
130-
"Failed to create .pgpass file. Please try manually with the "
131-
"create_pgpass_file() function."
132-
)
92+
self.connection = self.engine.connect()
93+
except Exception as err:
94+
if self._verbose:
95+
print(f"{err=}")
96+
self.engine = None
97+
if raise_err:
98+
raise err
99+
100+
def connect(self):
101+
"""Make a connection to the WRDS database."""
102+
# first try connection using system defaults and params set in constructor
103+
self.__make_sa_engine_conn()
104+
105+
if (self.engine is None and self._hostname != WRDS_POSTGRES_HOST):
106+
# try explicit w/ default hostname
107+
print(f"Trying '{WRDS_POSTGRES_HOST}'...")
108+
self._hostname = WRDS_POSTGRES_HOST
109+
self.__make_sa_engine_conn()
110+
111+
if (self.engine is None):
112+
# Use explicit username and password
113+
self._username, self._password = self.__get_user_credentials()
114+
# Last attempt, raise error if Exception encountered
115+
self.__make_sa_engine_conn(raise_err=True)
116+
117+
if (self.engine is None):
118+
print(f"Failed to connect {self._username}@{self._hostname}")
133119
else:
134-
print("You can create this file yourself at any time")
135-
print("with the create_pgpass_file() function.")
120+
# Connection successful. Offer to create a .pgpass for the user.
121+
print("WRDS recommends setting up a .pgpass file.")
122+
do_create_pgpass = ""
123+
while do_create_pgpass != "y" and do_create_pgpass != "n":
124+
do_create_pgpass = input("Create .pgpass file now [y/n]?: ")
125+
126+
if do_create_pgpass == "y":
127+
try:
128+
self.create_pgpass_file()
129+
print("Created .pgpass file successfully.")
130+
except:
131+
print("Failed to create .pgpass file.")
132+
print(
133+
"You can create this file yourself at any time "
134+
"with the create_pgpass_file() function."
135+
)
136136

137137
def close(self):
138138
"""
139139
Close the connection to the database.
140140
"""
141141
self.connection.close()
142142
self.engine.dispose()
143+
self.engine = None
143144

144145
def __enter__(self):
145146
self.connect()

0 commit comments

Comments
 (0)