first version of prober

Author: sfc-gh-pcyrek

cmd/prober/probes/__init__.py

@@ -1,6 +1,5 @@
 import logging
 
-
 def initialize_logger(name=__name__, level=logging.INFO):
     """
     Initializes and configures a logger.

cmd/__init__.py prober/__init__.pycmd/__init__.py renamed to prober/__init__.py

@@ -1,6 +1,5 @@
-from cmd.prober.logging_config import initialize_logger
-from cmd.prober.registry import prober_function
-
+from probes.logging_config import initialize_logger
+from probes.registry import prober_function
 import snowflake.connector
 
 # Initialize logger
@@ -30,6 +29,7 @@ def connect(connection_parameters):
             schema=connection_parameters.get("schema"),
             role=connection_parameters.get("role"),
             authenticator="KEY_PAIR_AUTHENTICATOR",
+            private_key=connection_parameters.get("private_key"),
         )
         return connection
     except Exception as e:

cmd/prober/__init__.py prober/probes/__init__.pycmd/prober/__init__.py renamed to prober/probes/__init__.py

@@ -1,26 +1,25 @@
 import argparse
 import logging
-from cmd.prober.logging_config import initialize_logger
-
-from .registry import PROBES_FUNCTIONS
+from probes.logging_config import initialize_logger
+from probes.registry import PROBES_FUNCTIONS
+import probes.login
 
 # Initialize logger
 logger = initialize_logger(__name__)
 
-
-if __name__ == "__main__":
+def main():
+    logger.info("Starting Python Driver Prober...")
     # Set up argument parser
     parser = argparse.ArgumentParser(description="Python Driver Prober")
     parser.add_argument("--scope", required=True, help="Scope of probing")
-
     parser.add_argument("--host", required=True, help="Host")
     parser.add_argument("--port", type=int, required=True, help="Port")
     parser.add_argument("--role", required=True, help="Protocol")
     parser.add_argument("--account", required=True, help="Account")
     parser.add_argument("--schema", required=True, help="Schema")
     parser.add_argument("--warehouse", required=True, help="Warehouse")
     parser.add_argument("--user", required=True, help="Username")
-    parser.add_argument("--password", required=True, help="Password")
+    parser.add_argument("--private_key", required=True, help="Private key")
 
     # Parse arguments
     args = parser.parse_args()
@@ -33,12 +32,17 @@
         "schema": args.schema,
         "warehouse": args.warehouse,
         "user": args.user,
-        "password": args.password,
+        "private_key": args.private_key,
     }
 
     for function_name, function in PROBES_FUNCTIONS.items():
         try:
+            logging.info("BBB")
             logging.error(f"Running probe: {function_name}")
             function(connection_params)
         except Exception as e:
             logging.error(f"Error running probe {function_name}: {e}")
+
+
+if __name__ == "__main__":
+    main()

cmd/prober/logging_config.py prober/probes/logging_config.pycmd/prober/logging_config.py renamed to prober/probes/logging_config.py

@@ -0,0 +1,16 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="snowflake_prober",
+    version="1.0.0",
+    packages=find_packages(),
+    install_requires=[
+        "snowflake-connector-python",
+        "requests",
+    ],
+    entry_points={
+        "console_scripts": [
+            "prober=probes.main:main",
+        ],
+    },
+)