Added class methods to initialize ES connection and some code refactoring

Author: vitaliyok

.gitignore

@@ -12,5 +12,5 @@ data/cogstack_search_results/
 # Default environments
 venv
 
-# pythin cache folder
+# python cache folder
 __pycache__

cogstack.py

@@ -21,9 +21,63 @@ class CogStack(object):
         hosts : List[str]
             A list of Elasticsearch host URLs.
     """
+    ES_TIMEOUT = 300
+    
     def __init__(self, hosts: List[str]):
         self.hosts = hosts
 
+    @classmethod
+    def with_basic_auth(cls, hosts: List[str], username: Optional[str] = None, password: Optional[str] = None) -> 'CogStack':
+        """
+        Create an instance of CogStack using basic authentication.
+
+        Parameters
+        ----------
+        hosts : List[str]
+            A list of Elasticsearch host URLs.
+        username : str, optional
+            The username to use when connecting to Elasticsearch. If not provided, the user will be prompted to enter a username.
+        password : str, optional
+            The password to use when connecting to Elasticsearch. If not provided, the user will be prompted to enter a password.
+        Returns 
+        -------
+            CogStack: An instance of the CogStack class.
+        """
+        cs = cls(hosts)
+        cs.use_basic_auth(username, password)
+        return cs
+    
+    @classmethod
+    def with_api_key_auth(cls, hosts: List[str], api_key: Optional[Dict] = None) -> 'CogStack':
+        """
+        Create an instance of CogStack using API key authentication.
+
+        Parameters
+        ----------
+        hosts : List[str]
+            A list of Elasticsearch host URLs.
+        apiKey : Dict, optional
+
+            API key object with "id" and "api_key" or "encoded" strings as fields. Generated in Elasticseach or Kibana
+            and provided by your CogStack administrator.
+            
+            If not provided, the user will be prompted to enter API key "encoded" value.
+            
+            Example:  
+                .. code-block:: json
+                        {
+                            "id": "API_KEY_ID",
+                            "api_key": "API_KEY",
+                            "encoded": "API_KEY_ENCODED_STRING"
+                        }
+        Returns
+        -------
+            CogStack: An instance of the CogStack class.
+        """
+        cs = cls(hosts)
+        cs.use_api_key_auth(api_key)
+        return cs
+
     def use_basic_auth(self, username: Optional[str] = None, password:Optional[str] = None) -> 'CogStack':
         """
         Create an instance of CogStack using basic authentication.
@@ -55,7 +109,7 @@ def use_api_key_auth(self,  api_key: Optional[Dict] = None) -> 'CogStack':
         ----------
         apiKey : Dict, optional
 
-            API key object with "id" and "api_key" or "encoded" strings as fields. Generated in Elasticseach or Kibana
+            API key object with "id" and "api_key" or "encoded" strings as fields. Generated in Elasticsearch or Kibana
             and provided by your CogStack administrator.
             
             If not provided, the user will be prompted to enter API key "encoded" value.
@@ -82,8 +136,14 @@ def use_api_key_auth(self,  api_key: Optional[Dict] = None) -> 'CogStack':
                 hasEncodedValue = True
             elif isinstance(api_key, Dict):
                 # If api_key is a dictionary, check for "encoded", "id" and "api_key" keys
-                encoded = api_key["encoded"] if "encoded" in api_key.keys() and api_key["encoded"] != '' else input("Encoded API key: ")
-                hasEncodedValue = encoded is not None and encoded != ''
+                if "id" in api_key.keys() and api_key["id"] != '' and "api_key" in api_key.keys() and api_key["api_key"] != '':
+                # If both "id" and "api_key" are present, use them
+                    encoded = None
+                    hasEncodedValue = False
+                else:
+                    # If "encoded" is present, use it; otherwise prompt for it
+                    encoded = api_key["encoded"] if "encoded" in api_key.keys() and api_key["encoded"] != '' else input("Encoded API key: ")
+                    hasEncodedValue = encoded is not None and encoded != ''
 
             if(not hasEncodedValue):
                 api_Id = api_key["id"] if "id" in api_key.keys() and api_key["id"] != '' else input("API Id: ")
@@ -110,7 +170,7 @@ def __connect(self, basic_auth : Optional[Tuple[str,str]] = None, api_key: Optio
                                                   api_key=api_key,
                                                   basic_auth=basic_auth,
                                                   verify_certs=False,
-                                                  request_timeout=300)
+                                                  request_timeout=self.ES_TIMEOUT)
         if not self.elastic.ping():
             raise Exception("CogStack connection failed. Please check your host list and credentials and try again.") 
         print("CogStack connection established successfully.")
@@ -202,7 +262,7 @@ def read_data_with_scan(self,
                             query: dict, 
                             include_fields: list[str]=None, 
                             size: int=1000, 
-                            request_timeout: int=300,
+                            request_timeout: int=ES_TIMEOUT,
                             show_progress: bool = True):
         """
         Retrieve documents from an Elasticsearch index using search query and elasticsearch scan helper function.
@@ -279,7 +339,7 @@ def read_data_with_scroll(self,
                   include_fields:Optional[list[str]]=None, 
                   size: int=1000, 
                   search_scroll_id: Optional[str] = None,
-                  request_timeout: Optional[int]=300,
+                  request_timeout: Optional[int]=ES_TIMEOUT,
                   show_progress: Optional[bool] = True):
 
         """
@@ -384,7 +444,7 @@ def read_data_with_sorting(self,
                                size: Optional[int]=1000, 
                                sort: Optional[dict|list[str]] = {"id": "asc"},
                                search_after: Optional[list[str|int|float|Any|None]] = None,
-                               request_timeout: Optional[int]=300,
+                               request_timeout: Optional[int]=ES_TIMEOUT,
                                show_progress: Optional[bool] = True):
         """
         Retrieve documents from an Elasticsearch index using search query and convert them to a Pandas DataFrame.

search/search_template.ipynb

@@ -35,7 +35,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "cs = CogStack(hosts).use_api_key_auth(api_key=api_key)"
+    "cs = CogStack.with_api_key_auth(hosts=hosts, api_key=api_key)\n",
+    "#cs = CogStack.with_basic_auth(hosts=hosts, username=username, password=password)\n",
+    "#cs = CogStack(hosts).use_api_key_auth(api_key=api_key)\n",
+    "#cs = CogStack(hosts).use_basic_auth(username=username, password=password)\n",
+    "#cs = CogStack(hosts).use_api_key_auth(\"\")\n"
    ]
   },
   {