Skip to content
This repository was archived by the owner on Mar 12, 2020. It is now read-only.

Commit 9c9c4cd

Browse files
tkopetstkopets
committed
Added new PostgreSQL auth method using env vars
Two new dict options under cli_options[type] are added in settings: * env - always sets the vars to given value * env_optional - sets the env vars only if sucessfully formatted Default settings file is updated to support this new functionality.
1 parent 90e8086 commit 9c9c4cd

File tree

3 files changed

+67
-17
lines changed

3 files changed

+67
-17
lines changed

‎SQLTools.sublime-settings

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,9 @@
6161
"options": ["--no-password"],
6262
"before": [],
6363
"args": "-h {host} -p {port} -U {username} -d {database}",
64+
"env_optional": {
65+
"PGPASSWORD": "{password}"
66+
},
6467
"queries": {
6568
"desc" : {
6669
"query": "select '|' || quote_ident(table_schema)||'.'||quote_ident(table_name) ||'|' as tblname from information_schema.tables where table_schema = any(current_schemas(false)) and table_schema not in ('pg_catalog', 'information_schema') order by table_schema = current_schema() desc, table_schema, table_name",

‎SQLToolsAPI/Command.py

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,14 @@
1010
class Command(object):
1111
timeout = 15
1212

13-
def __init__(self, args, callback, query=None, encoding='utf-8',
13+
def __init__(self, args, env, callback, query=None, encoding='utf-8',
1414
options=None, timeout=15, silenceErrors=False, stream=False):
1515
if options is None:
1616
options = {}
1717

1818
self.stream = stream
1919
self.args = args
20+
self.env = env
2021
self.callback = callback
2122
self.query = query
2223
self.encoding = encoding
@@ -43,13 +44,18 @@ def run(self):
4344
if self.silenceErrors:
4445
stderrHandle = subprocess.PIPE
4546

47+
# set the environment
48+
modifiedEnvironment = os.environ.copy()
49+
if (self.env):
50+
modifiedEnvironment.update(self.env)
51+
4652
queryTimerStart = time.time()
4753

4854
self.process = subprocess.Popen(self.args,
4955
stdout=subprocess.PIPE,
5056
stderr=stderrHandle,
5157
stdin=subprocess.PIPE,
52-
env=os.environ.copy(),
58+
env=modifiedEnvironment,
5359
startupinfo=si)
5460

5561
if self.stream:
@@ -103,21 +109,21 @@ def _formatShowQuery(query, queryTimeStart, queryTimeEnd):
103109
return resultString
104110

105111
@staticmethod
106-
def createAndRun(args, query, callback, options=None, timeout=15, silenceErrors=False, stream=False):
112+
def createAndRun(args, env, query, callback, options=None, timeout=15, silenceErrors=False, stream=False):
107113
if options is None:
108114
options = {}
109-
command = Command(args, callback, query, options=options,
115+
command = Command(args, env, callback, query, options=options,
110116
timeout=timeout, silenceErrors=silenceErrors)
111117
command.run()
112118

113119

114120
class ThreadCommand(Command, Thread):
115-
def __init__(self, args, callback, query=None, encoding='utf-8',
121+
def __init__(self, args, env, callback, query=None, encoding='utf-8',
116122
options=None, timeout=Command.timeout, silenceErrors=False, stream=False):
117123
if options is None:
118124
options = {}
119125

120-
Command.__init__(self, args, callback, query=query,
126+
Command.__init__(self, args, env, callback, query=query,
121127
encoding=encoding, options=options,
122128
timeout=timeout, silenceErrors=silenceErrors,
123129
stream=stream)
@@ -143,13 +149,13 @@ def stop(self):
143149
pass
144150

145151
@staticmethod
146-
def createAndRun(args, query, callback, options=None,
152+
def createAndRun(args, env, query, callback, options=None,
147153
timeout=Command.timeout, silenceErrors=False, stream=False):
148154
# Don't allow empty dicts or lists as defaults in method signature,
149155
# cfr http://nedbatchelder.com/blog/200806/pylint.html
150156
if options is None:
151157
options = {}
152-
command = ThreadCommand(args, callback, query, options=options,
158+
command = ThreadCommand(args, env, callback, query, options=options,
153159
timeout=timeout, silenceErrors=silenceErrors, stream=stream)
154160
command.start()
155161
killTimeout = Timer(command.timeout, command.stop)

‎SQLToolsAPI/Connection.py

Lines changed: 50 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -77,7 +77,9 @@ def getTables(self, callback):
7777
def cb(result):
7878
callback(U.getResultAsList(result))
7979

80-
self.Command.createAndRun(self.builArgs('desc'),
80+
args = self.buildArgs('desc')
81+
env = self.buildEnv()
82+
self.Command.createAndRun(args, env,
8183
query, cb, silenceErrors=True)
8284

8385
def getColumns(self, callback):
@@ -87,7 +89,9 @@ def cb(result):
8789

8890
try:
8991
query = self.getOptionsForSgdbCli()['queries']['columns']['query']
90-
self.Command.createAndRun(self.builArgs('columns'),
92+
args = self.buildArgs('columns')
93+
env = self.buildEnv()
94+
self.Command.createAndRun(args, env,
9195
query, cb, silenceErrors=True)
9296
except Exception:
9397
pass
@@ -99,26 +103,34 @@ def cb(result):
99103

100104
try:
101105
query = self.getOptionsForSgdbCli()['queries']['functions']['query']
102-
self.Command.createAndRun(self.builArgs('functions'),
106+
args = self.buildArgs('functions')
107+
env = self.buildEnv()
108+
self.Command.createAndRun(args, env,
103109
query, cb, silenceErrors=True)
104110
except Exception:
105111
pass
106112

107113
def getTableRecords(self, tableName, callback):
108114
query = self.getOptionsForSgdbCli()['queries']['show records']['query'].format(tableName, self.rowsLimit)
109115
queryToRun = '\n'.join(self.getOptionsForSgdbCli()['before'] + [query])
110-
self.Command.createAndRun(self.builArgs('show records'), queryToRun, callback, timeout=self.timeout)
116+
args = self.buildArgs('show records')
117+
env = self.buildEnv()
118+
self.Command.createAndRun(args, env, queryToRun, callback, timeout=self.timeout)
111119

112120
def getTableDescription(self, tableName, callback):
113121
query = self.getOptionsForSgdbCli()['queries']['desc table']['query'] % tableName
114122
queryToRun = '\n'.join(self.getOptionsForSgdbCli()['before'] + [query])
115-
self.Command.createAndRun(self.builArgs('desc table'), queryToRun, callback)
123+
args = self.buildArgs('desc table')
124+
env = self.buildEnv()
125+
self.Command.createAndRun(args, env, queryToRun, callback)
116126

117127
def getFunctionDescription(self, functionName, callback):
118128
query = self.getOptionsForSgdbCli()['queries']['desc function'][
119129
'query'] % functionName
120130
queryToRun = '\n'.join(self.getOptionsForSgdbCli()['before'] + [query])
121-
self.Command.createAndRun(self.builArgs('desc function'), queryToRun, callback)
131+
args = self.buildArgs('desc function')
132+
env = self.buildEnv()
133+
self.Command.createAndRun(args, env, queryToRun, callback)
122134

123135
def explainPlan(self, queries, callback):
124136
try:
@@ -132,7 +144,9 @@ def explainPlan(self, queries, callback):
132144
for query in filter(None, sqlparse.split(rawQuery))
133145
]
134146
queryToRun = '\n'.join(self.getOptionsForSgdbCli()['before'] + stripped_queries)
135-
self.Command.createAndRun(self.builArgs('explain plan'), queryToRun, callback, timeout=self.timeout)
147+
args = self.buildArgs('explain plan')
148+
env = self.buildEnv()
149+
self.Command.createAndRun(args, env, queryToRun, callback, timeout=self.timeout)
136150

137151
def execute(self, queries, callback, stream=False):
138152
queryToRun = ''
@@ -165,9 +179,11 @@ def execute(self, queries, callback, stream=False):
165179
if self.history:
166180
self.history.add(queryToRun)
167181

168-
self.Command.createAndRun(self.builArgs(), queryToRun, callback, options={'show_query': self.show_query}, timeout=self.timeout, stream=stream)
182+
args = self.buildArgs()
183+
env = self.buildEnv()
184+
self.Command.createAndRun(args, env, queryToRun, callback, options={'show_query': self.show_query}, timeout=self.timeout, stream=stream)
169185

170-
def builArgs(self, queryName=None):
186+
def buildArgs(self, queryName=None):
171187
cliOptions = self.getOptionsForSgdbCli()
172188
args = [self.cli]
173189

@@ -206,6 +222,31 @@ def builArgs(self, queryName=None):
206222
Log('Using cli args ' + ' '.join(args))
207223
return args
208224

225+
def buildEnv(self):
226+
cliOptions = self.getOptionsForSgdbCli()
227+
env = dict()
228+
229+
# append **optional** environment variables dict (if any)
230+
optionalEnv = cliOptions.get('env_optional')
231+
if optionalEnv: # only if we have optional args
232+
if isinstance(optionalEnv, dict):
233+
for var, value in optionalEnv.items():
234+
formattedValue = self.formatOptionalArgument(value, self.options)
235+
if formattedValue:
236+
env.update({var: formattedValue})
237+
238+
# append environment variables dict (if any)
239+
staticEnv = cliOptions.get('env')
240+
if staticEnv: # only if we have optional args
241+
if isinstance(staticEnv, dict):
242+
for var, value in staticEnv.items():
243+
formattedValue = value.format(**self.options)
244+
if formattedValue:
245+
env.update({var: formattedValue})
246+
247+
Log('Environment for command: ' + str(env))
248+
return env
249+
209250
def getOptionsForSgdbCli(self):
210251
return self.settings.get('cli_options', {}).get(self.type)
211252

0 commit comments

Comments
 (0)