def _init_rest(ref_sso_url, ref_proof_key, success=True, message=None):
def post_request(url, headers, body, **kwargs):
_ = url
_ = headers
_ = body
_ = kwargs.get('dummy')
return {
'success': success,
'message': message,
'data': {
'ssoUrl': ref_sso_url,
'proofKey': ref_proof_key,
}
}
connection = MagicMock()
connection._login_timeout = 120
connection.errorhandler = Mock(return_value=None)
type(connection).application = PropertyMock(return_value=CLIENT_NAME)
type(connection)._internal_application_name = PropertyMock(
return_value=CLIENT_NAME)
type(connection)._internal_application_version = PropertyMock(
return_value=CLIENT_VERSION)
rest = SnowflakeRestful(host='testaccount.snowflakecomputing.com',
port=443,
connection=connection)
rest._post_request = post_request
return rest
def _init_rest(ref_sso_url, ref_token_url, success=True, message=None):
def post_request(url, headers, body, **kwargs):
_ = url
_ = headers
_ = body
_ = kwargs.get("dummy")
return {
"success": success,
"message": message,
"data": {
"ssoUrl": ref_sso_url,
"tokenUrl": ref_token_url,
},
}
connection = MagicMock()
connection._login_timeout = 120
connection._network_timeout = None
connection.errorhandler = Mock(return_value=None)
connection._ocsp_mode = Mock(return_value=OCSPMode.FAIL_OPEN)
type(connection).application = PropertyMock(return_value=CLIENT_NAME)
type(connection)._internal_application_name = PropertyMock(
return_value=CLIENT_NAME)
type(connection)._internal_application_version = PropertyMock(
return_value=CLIENT_VERSION)
rest = SnowflakeRestful(host="testaccount.snowflakecomputing.com",
port=443,
connection=connection)
rest._post_request = post_request
return rest
def _init_rest(application, post_requset):
connection = MagicMock()
connection._login_timeout = 120
connection.errorhandler = Mock(return_value=None)
type(connection).application = PropertyMock(return_value=application)
type(connection)._internal_application_name = PropertyMock(
return_value=CLIENT_NAME)
type(connection)._internal_application_version = PropertyMock(
return_value=CLIENT_VERSION)
rest = SnowflakeRestful(host='testaccount.snowflakecomputing.com',
port=443,
connection=connection)
rest._post_request = post_requset
return rest
def __open_connection(self):
u"""
Opens a new network connection
"""
self.converter = self._converter_class(use_sfbinaryformat=False,
use_numpy=self._numpy)
self._rest = SnowflakeRestful(
host=self.host,
port=self.port,
protocol=self._protocol,
inject_client_pause=self._inject_client_pause,
connection=self)
if self.host.endswith(u".privatelink.snowflakecomputing.com"):
ocsp_cache_server = \
u'http://ocsp{}/ocsp_response_cache.json'.format(
self.host[self.host.index('.'):])
os.environ[
'SF_OCSP_RESPONSE_CACHE_SERVER_URL'] = ocsp_cache_server
else:
if 'SF_OCSP_RESPONSE_CACHE_SERVER_URL' in os.environ:
del os.environ['SF_OCSP_RESPONSE_CACHE_SERVER_URL']
auth_instance = AuthByWebBrowser(self.rest,
self.application,
protocol=self._protocol,
host=self.host,
port=self.port,
webbrowser_pkg=WebbrowserPkg,
socket_pkg=SocketPkg)
if self._session_parameters is None:
self._session_parameters = {}
if self._autocommit is not None:
self._session_parameters['AUTOCOMMIT'] = self._autocommit
if self._timezone is not None:
self._session_parameters['TIMEZONE'] = self._timezone
if self.client_session_keep_alive:
self._session_parameters['CLIENT_SESSION_KEEP_ALIVE'] = True
# enable storing temporary credential in a file
self._session_parameters[
'CLIENT_STORE_TEMPORARY_CREDENTIAL'] = True
auth = Auth(self.rest)
if not auth.read_temporary_credential(self.account, self.user,
self._session_parameters):
self.__authenticate(auth_instance)
else:
# set the current objects as the session is derived from the id
# token, and the current objects may be different.
self._set_current_objects()
self._password = None # ensure password won't persist
if self.client_session_keep_alive:
self._add_heartbeat()
def create_session(rest: SnowflakeRestful,
num_sessions: int = 1,
url: str | None = None) -> None:
"""
Creates 'num_sessions' sessions to 'url'. This is recursive so that idle sessions
are not reused.
"""
if num_sessions == 0:
return
with rest._use_requests_session(url):
create_session(rest, num_sessions - 1, url)
def test_no_auth(db_parameters):
"""
SNOW-13588: No auth Rest API test
"""
rest = SnowflakeRestful(
host=db_parameters['host'],
port=db_parameters['port'])
try:
# no auth
# show warehouse
rest.request(
url='/queries',
body={
'sequenceId': 10000,
'sqlText': 'show warehouses',
'parameters': {
'ui_mode': True,
},
},
method='post',
client='rest')
raise Exception("Must fail with auth error")
except errors.Error as e:
assert e.errno == errorcode.ER_CONNECTION_IS_CLOSED
finally:
rest.close()
def test_no_url_multiple_sessions(make_session_mock):
rest = SnowflakeRestful(connection=mock_conn)
create_session(rest, 2)
assert make_session_mock.call_count == 2
assert list(rest._sessions_map.keys()) == [None]
session_pool = rest._sessions_map[None]
assert len(session_pool._idle_sessions) == 2
assert len(session_pool._active_sessions) == 0
close_sessions(rest, 1)
def test_multiple_urls_multiple_sessions(make_session_mock):
rest = SnowflakeRestful(connection=mock_conn)
for url in [url_1, url_2, None]:
create_session(rest, num_sessions=2, url=url)
assert make_session_mock.call_count == 6
hostnames = list(rest._sessions_map.keys())
for hostname in [hostname_1, hostname_2, None]:
assert hostname in hostnames
for pool in rest._sessions_map.values():
assert len(pool._idle_sessions) == 2
assert len(pool._active_sessions) == 0
close_sessions(rest, 3)
def test_multiple_urls_reuse_sessions(make_session_mock):
rest = SnowflakeRestful(connection=mock_conn)
for url in [url_1, url_2, url_3, None]:
# create 10 sessions, one after another
for _ in range(10):
create_session(rest, url=url)
# only one session is created and reused thereafter
assert make_session_mock.call_count == 3
hostnames = list(rest._sessions_map.keys())
assert len(hostnames) == 3
for hostname in [hostname_1, hostname_2, None]:
assert hostname in hostnames
for pool in rest._sessions_map.values():
assert len(pool._idle_sessions) == 1
assert len(pool._active_sessions) == 0
close_sessions(rest, 3)
def test_no_auth(db_parameters):
"""SNOW-13588: No auth Rest API test."""
rest = SnowflakeRestful(host=db_parameters["host"],
port=db_parameters["port"])
try:
# no auth
# show warehouse
rest.request(
url="/queries",
body={
"sequenceId": 10000,
"sqlText": "show warehouses",
"parameters": {
"ui_mode": True,
},
},
method="post",
client="rest",
)
raise Exception("Must fail with auth error")
except errors.Error as e:
assert e.errno == errorcode.ER_CONNECTION_IS_CLOSED
finally:
rest.close()
def test_request_exec():
rest = SnowflakeRestful(host="testaccount.snowflakecomputing.com",
port=443)
default_parameters = {
"method": "POST",
"full_url": "https://testaccount.snowflakecomputing.com/",
"headers": {},
"data": '{"code": 12345}',
"token": None,
}
# request mock
output_data = {"success": True, "code": 12345}
request_mock = MagicMock()
type(request_mock).status_code = PropertyMock(return_value=OK)
request_mock.json.return_value = output_data
# session mock
session = MagicMock()
session.request.return_value = request_mock
# success
ret = rest._request_exec(session=session, **default_parameters)
assert ret == output_data, "output data"
# retryable exceptions
for errcode in [
BAD_REQUEST, # 400
FORBIDDEN, # 403
INTERNAL_SERVER_ERROR, # 500
BAD_GATEWAY, # 502
SERVICE_UNAVAILABLE, # 503
GATEWAY_TIMEOUT, # 504
555, # random 5xx error
]:
type(request_mock).status_code = PropertyMock(return_value=errcode)
try:
rest._request_exec(session=session, **default_parameters)
pytest.fail("should fail")
except RetryRequest as e:
cls = STATUS_TO_EXCEPTION.get(errcode, OtherHTTPRetryableError)
assert isinstance(e.args[0],
cls), "must be internal error exception"
# unauthorized
type(request_mock).status_code = PropertyMock(return_value=UNAUTHORIZED)
with pytest.raises(InterfaceError):
rest._request_exec(session=session, **default_parameters)
# unauthorized with catch okta unauthorized error
# TODO: what is the difference to InterfaceError?
type(request_mock).status_code = PropertyMock(return_value=UNAUTHORIZED)
with pytest.raises(DatabaseError):
rest._request_exec(session=session,
catch_okta_unauthorized_error=True,
**default_parameters)
class IncompleteReadMock(IncompleteRead):
def __init__(self):
IncompleteRead.__init__(self, "")
# handle retryable exception
for exc in [
requests.exceptions.ConnectTimeout,
requests.exceptions.ReadTimeout,
IncompleteReadMock,
urllib3.exceptions.ProtocolError,
requests.exceptions.ConnectionError,
AttributeError,
]:
session = MagicMock()
session.request = Mock(side_effect=exc)
try:
rest._request_exec(session=session, **default_parameters)
pytest.fail("should fail")
except RetryRequest as e:
cause = e.args[0]
assert isinstance(cause, exc), "same error class"
# handle OpenSSL errors and BadStateLine
for exc in [
OpenSSL.SSL.SysCallError(errno.ECONNRESET),
OpenSSL.SSL.SysCallError(errno.ETIMEDOUT),
OpenSSL.SSL.SysCallError(errno.EPIPE),
OpenSSL.SSL.SysCallError(-1), # unknown
# TODO: should we keep this?
# urllib3.exceptions.ReadTimeoutError(None, None, None),
BadStatusLine("fake"),
]:
session = MagicMock()
session.request = Mock(side_effect=exc)
try:
rest._request_exec(session=session, **default_parameters)
pytest.fail("should fail")
except RetryRequest as e:
assert e.args[0] == exc, "same error instance"
def test_fetch():
connection = MagicMock()
connection.errorhandler = Mock(return_value=None)
rest = SnowflakeRestful(host="testaccount.snowflakecomputing.com",
port=443,
connection=connection)
class Cnt:
def __init__(self):
self.c = 0
def set(self, cnt):
self.c = cnt
def reset(self):
self.set(0)
cnt = Cnt()
default_parameters = {
"method": "POST",
"full_url": "https://testaccount.snowflakecomputing.com/",
"headers": {
"cnt": cnt
},
"data": '{"code": 12345}',
}
NOT_RETRYABLE = 1000
class NotRetryableException(Exception):
pass
def fake_request_exec(**kwargs):
headers = kwargs.get("headers")
cnt = headers["cnt"]
time.sleep(3)
if cnt.c <= 1:
# the first two raises failure
cnt.c += 1
raise RetryRequest(Exception("can retry"))
elif cnt.c == NOT_RETRYABLE:
# not retryable exception
raise NotRetryableException("cannot retry")
else:
# return success in the third attempt
return {"success": True, "data": "valid data"}
# inject a fake method
rest._request_exec = fake_request_exec
# first two attempts will fail but third will success
cnt.reset()
ret = rest.fetch(timeout=10, **default_parameters)
assert ret == {"success": True, "data": "valid data"}
assert not rest._connection.errorhandler.called # no error
# first attempt to reach timeout even if the exception is retryable
cnt.reset()
ret = rest.fetch(timeout=1, **default_parameters)
assert ret == {}
assert rest._connection.errorhandler.called # error
# not retryable excpetion
cnt.set(NOT_RETRYABLE)
with pytest.raises(NotRetryableException):
rest.fetch(timeout=7, **default_parameters)
# first attempt fails and will not retry
cnt.reset()
default_parameters["no_retry"] = True
ret = rest.fetch(timeout=10, **default_parameters)
assert ret == {}
assert cnt.c == 1 # failed on first call - did not retry
assert rest._connection.errorhandler.called # error
def test_request_exec():
rest = SnowflakeRestful(
host='testaccount.snowflakecomputing.com',
port=443)
default_parameters = {
'method': "POST",
'full_url': "https://testaccount.snowflakecomputing.com/",
'headers': {},
'data': '{"code": 12345}',
'proxies': None,
'token': None
}
# request mock
output_data = {'success': True, 'code': 12345}
request_mock = MagicMock()
type(request_mock).status_code = PropertyMock(return_value=OK)
request_mock.json.return_value = output_data
# session mock
session = MagicMock()
session.request.return_value = request_mock
# success
ret = rest._request_exec(session=session, **default_parameters)
assert ret == output_data, 'output data'
# retryable exceptions
for errcode in [
BAD_REQUEST, # 400
FORBIDDEN, # 403
INTERNAL_SERVER_ERROR, # 500
BAD_GATEWAY, # 502
SERVICE_UNAVAILABLE, # 503
GATEWAY_TIMEOUT, # 504
555, # random 5xx error
]:
type(request_mock).status_code = PropertyMock(return_value=errcode)
try:
rest._request_exec(
session=session, **default_parameters)
pytest.fail('should fail')
except RetryRequest as e:
cls = STATUS_TO_EXCEPTION.get(
errcode,
OtherHTTPRetryableError)
assert isinstance(
e.args[0],
cls), "must be internal error exception"
# unauthorized
type(request_mock).status_code = PropertyMock(return_value=UNAUTHORIZED)
with pytest.raises(InterfaceError):
rest._request_exec(
session=session, **default_parameters)
# unauthorized with catch okta unauthorized error
# TODO: what is the difference to InterfaceError?
type(request_mock).status_code = PropertyMock(return_value=UNAUTHORIZED)
with pytest.raises(DatabaseError):
rest._request_exec(
session=session, catch_okta_unauthorized_error=True,
**default_parameters)
# handle retryable exception
for exc in [
ConnectTimeout,
ReadTimeout,
SSLError,
ProtocolError,
ConnectionError,
AttributeError,
PyAsn1Error
]:
session = MagicMock()
session.request = Mock(side_effect=exc)
try:
rest._request_exec(
session=session, **default_parameters)
pytest.fail('should fail')
except RetryRequest as e:
cause = e.args[0]
assert isinstance(cause, exc), "same error class"
# handle OpenSSL errors and BadStateLine
for exc in [
OpenSSL.SSL.SysCallError(errno.ECONNRESET),
OpenSSL.SSL.SysCallError(errno.ETIMEDOUT),
OpenSSL.SSL.SysCallError(errno.EPIPE),
OpenSSL.SSL.SysCallError(-1), # unknown
ReadTimeoutError(None, None, None),
BadStatusLine('fake')
]:
session = MagicMock()
session.request = Mock(side_effect=exc)
try:
rest._request_exec(
session=session, **default_parameters)
pytest.fail('should fail')
except RetryRequest as e:
assert e.args[0] == exc, "same error instance"
def test_fetch():
connection = MagicMock()
connection.errorhandler = Mock(return_value=None)
rest = SnowflakeRestful(
host='testaccount.snowflakecomputing.com',
port=443,
connection=connection)
class Cnt(object):
def __init__(self):
self.c = 0
def set(self, cnt):
self.c = cnt
def reset(self):
self.set(0)
cnt = Cnt()
default_parameters = {
'method': "POST",
'full_url': "https://testaccount.snowflakecomputing.com/",
'headers': {'cnt': cnt},
'data': '{"code": 12345}',
}
NOT_RETRYABLE = 1000
class NotRetryableException(Exception):
pass
def fake_request_exec(**kwargs):
headers = kwargs.get('headers')
cnt = headers['cnt']
time.sleep(3)
if cnt.c <= 1:
# the first two raises failure
cnt.c += 1
raise RetryRequest(Exception('can retry'))
elif cnt.c == NOT_RETRYABLE:
# not retryable exception
raise NotRetryableException('cannot retry')
else:
# return success in the third attempt
return {'success': True, 'data': "valid data"}
# inject a fake method
rest._request_exec = fake_request_exec
# first two attempts will fail but third will success
cnt.reset()
ret = rest.fetch(timeout=10, **default_parameters)
assert ret == {'success': True, 'data': "valid data"}
assert not rest._connection.errorhandler.called # no error
# first attempt to reach timeout even if the exception is retryable
cnt.reset()
ret = rest.fetch(timeout=1, **default_parameters)
assert ret == {}
assert rest._connection.errorhandler.called # error
# not retryable excpetion
cnt.set(NOT_RETRYABLE)
with pytest.raises(NotRetryableException):
rest.fetch(timeout=7, **default_parameters)
def test_renew_session():
OLD_SESSION_TOKEN = "old_session_token"
OLD_MASTER_TOKEN = "old_master_token"
NEW_SESSION_TOKEN = "new_session_token"
NEW_MASTER_TOKEN = "new_master_token"
connection = MagicMock()
connection.errorhandler = Mock(return_value=None)
type(connection)._probe_connection = PropertyMock(return_value=False)
rest = SnowflakeRestful(host='testaccount.snowflakecomputing.com',
port=443,
connection=connection)
rest._token = OLD_SESSION_TOKEN
rest._master_token = OLD_MASTER_TOKEN
# inject a fake method (success)
def fake_request_exec(**_):
return {
u'success': True,
u'data': {
u"sessionToken": NEW_SESSION_TOKEN,
u"masterToken": NEW_MASTER_TOKEN
}
}
rest._request_exec = fake_request_exec
rest._renew_session()
assert not rest._connection.errorhandler.called # no error
assert rest.master_token == NEW_MASTER_TOKEN
assert rest.token == NEW_SESSION_TOKEN
# inject a fake method (failure)
def fake_request_exec(**_):
return {
u'success': False,
u'message': "failed to renew session",
u'code': 987654
}
rest._request_exec = fake_request_exec
rest._renew_session()
assert rest._connection.errorhandler.called # error
# no master token
del rest._master_token
rest._renew_session()
assert rest._connection.errorhandler.called # error
def close_sessions(rest: SnowflakeRestful, num_session_pools: int) -> None:
"""Helper function to call SnowflakeRestful.close(). Asserts close was called on all SessionPools."""
with mock.patch(
"snowflake.connector.network.SessionPool.close") as close_mock:
rest.close()
assert close_mock.call_count == num_session_pools
