Merge pull request #168 from mindsdb/feat/add-cookies

feat: add cookies to requests

Author: lucas-koontz

mindsdb_sdk/__about__.py

@@ -1,6 +1,6 @@
 __title__ = 'mindsdb_sdk'
 __package_name__ = 'mindsdb_sdk'
-__version__ = '3.1.6'
+__version__ = '3.1.7'
 __description__ = "MindsDB Python SDK, provides an SDK to use a remote mindsdb instance"
 __email__ = "[email protected]"
 __author__ = 'MindsDB Inc'

mindsdb_sdk/connect.py

@@ -12,6 +12,7 @@ def connect(
         password: str = None,
         api_key: str = None,
         is_managed: bool = False,
+        cookies=None,
         headers=None) -> Server:
     """
     Create connection to mindsdb server
@@ -21,6 +22,7 @@ def connect(
     :param password: user password to login (for cloud version)
     :param api_key: API key to authenticate (for cloud version)
     :param is_managed: whether or not the URL points to a managed instance
+    :param cookies: addtional cookies to send with the connection, optional
     :param headers: addtional headers to send with the connection, optional
     :return: Server object
 
@@ -51,6 +53,7 @@ def connect(
             # is local
             url = DEFAULT_LOCAL_API_URL
 
-    api = RestAPI(url, login, password, api_key, is_managed, headers=headers)
+    api = RestAPI(url, login, password, api_key, is_managed, 
+                  cookies=cookies, headers=headers)
 
     return Server(api)

mindsdb_sdk/connectors/rest_api.py

@@ -37,7 +37,8 @@ def _raise_for_status(response):
 
 
 class RestAPI:
-    def __init__(self, url=None, login=None, password=None, api_key=None, is_managed=False, headers=None):
+    def __init__(self, url=None, login=None, password=None, api_key=None, is_managed=False, 
+                 cookies=None, headers=None):
 
         self.url = url
         self.username = login
@@ -46,6 +47,9 @@ def __init__(self, url=None, login=None, password=None, api_key=None, is_managed
         self.is_managed = is_managed
         self.session = requests.Session()
 
+        if cookies is not None:
+            self.session.cookies.update(cookies)
+
         self.session.headers['User-Agent'] = f'python-sdk/{__about__.__version__}'
         if headers is not None:
             self.session.headers.update(headers)