kubernetes-watch 0.1.5__py3-none-any.whl → 0.1.9__py3-none-any.whl

kube_watch/modules/database/postgre.py

@@ -0,0 +1,271 @@
+import psycopg2
+import psycopg2.extras
+from prefect import get_run_logger
+
+from .model import TableQuery
+
+logger = get_run_logger()
+
+
+def execute_query(db_user, db_pass, db_query, db_host="localhost", db_port=5432, db_name="postgres"):
+    """
+    Connect to PostgreSQL database, execute a query, and return status message.
+    
+    Args:
+        db_user (str): Database username
+        db_pass (str): Database password
+        db_query (str): SQL query to execute
+        db_host (str): Database host (default: localhost)
+        db_port (int): Database port (default: 5432)
+        db_name (str): Database name (default: postgres)
+    
+    Returns:
+        dict: Status message with success/failure information
+    """
+    connection = None
+    cursor = None
+    
+    try:
+        # Establish database connection
+        connection = psycopg2.connect(
+            host=db_host,
+            port=db_port,
+            database=db_name,
+            user=db_user,
+            password=db_pass
+        )
+        
+        cursor = connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
+        
+        # Execute the query
+        cursor.execute(db_query)
+        
+        # Check if it's a SELECT query to fetch results
+        if db_query.strip().upper().startswith('SELECT'):
+            results = cursor.fetchall()
+            row_count = len(results)
+            connection.commit()  # Commit even for SELECT to close transaction
+            
+            logger.info(f"Query executed successfully. Retrieved {row_count} rows.")
+            return {
+                "status": "success",
+                "message": f"Query executed successfully. Retrieved {row_count} rows.",
+                "row_count": row_count,
+                "data": results
+            }
+        else:
+            # For INSERT, UPDATE, DELETE queries
+            connection.commit()
+            affected_rows = cursor.rowcount
+            
+            logger.info(f"Query executed successfully. {affected_rows} rows affected.")
+            return {
+                "status": "success",
+                "message": f"Query executed successfully. {affected_rows} rows affected.",
+                "affected_rows": affected_rows
+            }
+            
+    except psycopg2.Error as e:
+        # PostgreSQL specific errors
+        error_msg = f"PostgreSQL error: {str(e)}"
+        logger.error(error_msg)
+        
+        if connection:
+            connection.rollback()
+            
+        return {
+            "status": "error",
+            "message": error_msg,
+            "error_type": "postgresql_error"
+        }
+        
+    except Exception as e:
+        # General errors
+        error_msg = f"Unexpected error: {str(e)}"
+        logger.error(error_msg)
+        
+        if connection:
+            connection.rollback()
+            
+        return {
+            "status": "error", 
+            "message": error_msg,
+            "error_type": "general_error"
+        }
+        
+    finally:
+        # Clean up connections
+        try:
+            if cursor:
+                cursor.close()
+                logger.debug("Database cursor closed.")
+                
+            if connection:
+                connection.close()
+                logger.debug("Database connection closed.")
+                
+        except Exception as cleanup_error:
+            logger.warning(f"Error during cleanup: {str(cleanup_error)}")
+
+
+def delete_on_retention_period(table_delete: dict, batch_size: int = 100000, interval_days: int = 14):
+    """
+    Delete rows older than a specified retention period from a table in PostgreSQL.
+
+    Args:
+        table_delete (dict[TableQuery]): Object containing table name and column name.
+        batch_size (int): Number of rows to delete in each batch (default: 100000).
+        interval_days (int): Retention period in days (default: 14).
+
+    Returns:
+        dict: Status message with success/failure information.
+    """
+
+    try:
+        table_query = TableQuery(**table_delete)
+    except Exception as e:
+        logger.error(f"Error creating TableQuery object: {str(e)}")
+        raise ValueError("Invalid table_delete data format. Expected a dictionary with 'name', 'column_name', 'db_host', 'db_port', 'db_name', 'db_user', and 'db_pass' keys.")
+
+    connection = None
+    cursor = None
+
+    try:
+        # Establish database connection
+        connection = psycopg2.connect(
+            host=table_query.db_host,
+            port=table_query.db_port,
+            database=table_query.db_name,
+            user=table_query.db_user,
+            password=table_query.db_pass
+        )
+
+        cursor = connection.cursor()
+
+        rows_deleted_total = 0
+
+        while True:
+            # Build the DELETE query dynamically
+            delete_query = f"""
+            DELETE FROM {table_query.table_name}
+            WHERE id IN (
+                SELECT id FROM {table_query.table_name}
+                WHERE {table_query.column_name} < NOW() - INTERVAL '{interval_days} days'
+                LIMIT {batch_size}
+            )
+            """
+
+            cursor.execute(delete_query)
+            rows_deleted = cursor.rowcount
+            rows_deleted_total += rows_deleted
+
+            connection.commit()
+
+            logger.info(f"Deleted {rows_deleted} rows from {table_query.table_name}.")
+
+            # Exit the loop if no rows were deleted
+            if rows_deleted == 0:
+                break
+
+        logger.info(f"Total rows deleted from {table_query.table_name}: {rows_deleted_total}.")
+
+        return {
+            "status": "success",
+            "message": f"Total rows deleted from {table_query.table_name}: {rows_deleted_total}.",
+            "rows_deleted_total": rows_deleted_total
+        }
+
+    except psycopg2.Error as e:
+        error_msg = f"PostgreSQL error: {str(e)}"
+        logger.error(error_msg)
+
+        if connection:
+            connection.rollback()
+
+        return {
+            "status": "error",
+            "message": error_msg,
+            "error_type": "postgresql_error"
+        }
+
+    except Exception as e:
+        error_msg = f"Unexpected error: {str(e)}"
+        logger.error(error_msg)
+
+        if connection:
+            connection.rollback()
+
+        return {
+            "status": "error",
+            "message": error_msg,
+            "error_type": "general_error"
+        }
+
+    finally:
+        # Clean up connections
+        try:
+            if cursor:
+                cursor.close()
+                logger.debug("Database cursor closed.")
+
+            if connection:
+                connection.close()
+                logger.debug("Database connection closed.")
+
+        except Exception as cleanup_error:
+            logger.warning(f"Error during cleanup: {str(cleanup_error)}")
+
+"""
+EXAMPLE RETENTION PERIOD DELETION QUERY
+
+DO $$
+DECLARE
+    batch_size INTEGER := 100000;
+    rows_deleted INTEGER;
+    t TEXT;
+    col TEXT;
+    sql TEXT;
+    count_sql TEXT;
+
+    -- Define loop record
+    rec RECORD;
+BEGIN
+    -- Simulate table/column pairs using VALUES
+    FOR rec IN
+        SELECT * FROM (
+            VALUES
+                ('log', 'created'),
+                ('task_run_state', 'timestamp'),
+                ('task_run', 'start_time'),
+                ('flow_run_state', 'timestamp'),
+                ('flow_run', 'start_time')
+        ) AS table_info(table_name, column_name)
+    LOOP
+        t := rec.table_name;
+        col := rec.column_name;
+
+        LOOP
+            sql := format(
+                'DELETE FROM %I WHERE id IN (
+                    SELECT id FROM %I WHERE %I < NOW() - INTERVAL ''14 days'' LIMIT %s
+                )',
+                t, t, col, batch_size
+            );
+
+            EXECUTE sql;
+            GET DIAGNOSTICS rows_deleted = ROW_COUNT;
+
+            RAISE NOTICE 'Deleted % rows from %', rows_deleted, t;
+            EXIT WHEN rows_deleted = 0;
+        END LOOP;
+
+        count_sql := format(
+            'SELECT COUNT(*) FROM %I WHERE %I < NOW() - INTERVAL ''14 days''',
+            t, col
+        );
+
+        EXECUTE count_sql INTO rows_deleted;
+        RAISE NOTICE 'Remaining rows in % older than 14 days: %', t, rows_deleted;
+    END LOOP;
+END $$;
+"""

kube_watch/modules/logic/actions.py

@@ -1,56 +1,56 @@
-import subprocess
-import os
-from typing import List
-from prefect import get_run_logger
-logger = get_run_logger()
-
-def run_standalone_script(package_name, package_run, package_exec):
-    script_dir = os.path.dirname(os.path.realpath(__file__))
-    # script_path = os.path.join(script_dir, package_name.replace('.', os.sep))
-    target_dir = os.path.join(script_dir, os.pardir, os.pardir, *package_name.split('.'))
-
-    full_command = f"{package_run} {os.path.join(target_dir, package_exec)}"
-
-    # Execute the command
-    try:
-        result = subprocess.run(full_command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
-        if result.stdout:
-            logger.info(result.stdout)
-        if result.stderr:
-            logger.error(result.stderr)
-        # logger.info(f"Output: {result.stdout}")
-        result.check_returncode()
-    except subprocess.CalledProcessError as e:
-        # All logs should have already been handled above, now just raise an exception
-        logger.error("The subprocess encountered an error: %s", e)
-        raise Exception("Subprocess failed with exit code {}".format(e.returncode))
-    
-
-def run_standalone_script_modified(base_path: str, package_name: str, package_run_cmds: List[str]):
-    # Construct the absolute path to the target directory
-    target_dir = os.path.join(base_path, *package_name.split('.'))
-
-    commands = [f"cd {target_dir}"] + package_run_cmds
-    full_command = " && ".join(commands)
-
-    # full_command = f"cd {target_dir} && {package_run_cmd}"
-
-    # Build the full command to execute
-    # full_command = f"{package_run} {os.path.join(target_dir, package_exec)}"
-
-    # print(full_command)
-
-    # Execute the command
-    try:
-        result = subprocess.run(full_command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
-        if result.stdout:
-            logger.info(result.stdout)
-        if result.stderr:
-            logger.error(result.stderr)
-        result.check_returncode()
-    except subprocess.CalledProcessError as e:
-        logger.error("Command failed with exit code %s", e.returncode)
-        logger.error("Output:\n%s", e.stdout)
-        logger.error("Errors:\n%s", e.stderr)
-        raise Exception(f"Subprocess failed with exit code {e.returncode}. Check logs for more details.")
+import subprocess
+import os
+from typing import List
+from prefect import get_run_logger
+logger = get_run_logger()
+
+def run_standalone_script(package_name, package_run, package_exec):
+    script_dir = os.path.dirname(os.path.realpath(__file__))
+    # script_path = os.path.join(script_dir, package_name.replace('.', os.sep))
+    target_dir = os.path.join(script_dir, os.pardir, os.pardir, *package_name.split('.'))
+
+    full_command = f"{package_run} {os.path.join(target_dir, package_exec)}"
+
+    # Execute the command
+    try:
+        result = subprocess.run(full_command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+        if result.stdout:
+            logger.info(result.stdout)
+        if result.stderr:
+            logger.error(result.stderr)
+        # logger.info(f"Output: {result.stdout}")
+        result.check_returncode()
+    except subprocess.CalledProcessError as e:
+        # All logs should have already been handled above, now just raise an exception
+        logger.error("The subprocess encountered an error: %s", e)
+        raise Exception("Subprocess failed with exit code {}".format(e.returncode))
+    
+
+def run_standalone_script_modified(base_path: str, package_name: str, package_run_cmds: List[str]):
+    # Construct the absolute path to the target directory
+    target_dir = os.path.join(base_path, *package_name.split('.'))
+
+    commands = [f"cd {target_dir}"] + package_run_cmds
+    full_command = " && ".join(commands)
+
+    # full_command = f"cd {target_dir} && {package_run_cmd}"
+
+    # Build the full command to execute
+    # full_command = f"{package_run} {os.path.join(target_dir, package_exec)}"
+
+    # print(full_command)
+
+    # Execute the command
+    try:
+        result = subprocess.run(full_command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+        if result.stdout:
+            logger.info(result.stdout)
+        if result.stderr:
+            logger.error(result.stderr)
+        result.check_returncode()
+    except subprocess.CalledProcessError as e:
+        logger.error("Command failed with exit code %s", e.returncode)
+        logger.error("Output:\n%s", e.stdout)
+        logger.error("Errors:\n%s", e.stderr)
+        raise Exception(f"Subprocess failed with exit code {e.returncode}. Check logs for more details.")
         # raise Exception(f"Subprocess failed with exit code {e.returncode}")

kube_watch/modules/logic/checks.py

@@ -1,8 +1,8 @@
-
-
-def dicts_has_diff(dict_a, dict_b):
-    return dict_a != dict_b
-
-
-def remove_keys(d, keys):
+
+
+def dicts_has_diff(dict_a, dict_b):
+    return dict_a != dict_b
+
+
+def remove_keys(d, keys):
     return {k: v for k, v in d.items() if k not in keys}

kube_watch/modules/logic/load.py

@@ -1,23 +1,23 @@
-import os
-from prefect import get_run_logger
-logger = get_run_logger()
-
-def load_secrets_to_env(data):
-    for key, value in data.items():
-        if key in os.environ:
-            del os.environ[key]
-        os.environ[key] = value
-        # logger.info(f"ENV VAR: {key} loaded")
-
-def load_env_from_file(filepath):
-    with open(filepath, "r") as f:
-        for line in f:
-            # Remove whitespace and ignore comments
-            line = line.strip()
-            if line and not line.startswith('#'):
-                key, value = line.split('=', 1)
-                # Remove the environment variable if it already exists
-                if key in os.environ:
-                    del os.environ[key]
-                # Set the new value
-                os.environ[key] = value
+import os
+from prefect import get_run_logger
+logger = get_run_logger()
+
+def load_secrets_to_env(data):
+    for key, value in data.items():
+        if key in os.environ:
+            del os.environ[key]
+        os.environ[key] = value
+        # logger.info(f"ENV VAR: {key} loaded")
+
+def load_env_from_file(filepath):
+    with open(filepath, "r") as f:
+        for line in f:
+            # Remove whitespace and ignore comments
+            line = line.strip()
+            if line and not line.startswith('#'):
+                key, value = line.split('=', 1)
+                # Remove the environment variable if it already exists
+                if key in os.environ:
+                    del os.environ[key]
+                # Set the new value
+                os.environ[key] = value

kube_watch/modules/logic/merge.py

@@ -1,31 +1,31 @@
-from typing import Any, List, Dict
-from kube_watch.enums.logic import Operations
-
-
-def merge_logical_outputs(inp_dict: Dict):
-    if 'operation' not in inp_dict.keys():
-        raise TypeError("Missing required parameters: 'operation'")
-    operation = inp_dict.get('operation')
-    del inp_dict['operation']
-
-    inputs = [v for k,v in inp_dict.items()]
-    return merge_logical_list(inputs, operation)
-
-
-def merge_logical_list(inp_list: List, operation: Operations):
-    if operation == Operations.OR:
-        return any(inp_list)
-    if operation == Operations.AND:
-        return all(inp_list)
-    raise ValueError("Invalid logical operation")
-
-
-def partial_dict_update(orig_data, new_data):
-    """
-    This function is used when some key value pairs in orig_data should
-    be updated from new_data.
-    """
-    for k, v in new_data.items():
-        orig_data[k] = v
-
-    return orig_data
+from typing import Any, List, Dict
+from kube_watch.enums.logic import Operations
+
+
+def merge_logical_outputs(inp_dict: Dict):
+    if 'operation' not in inp_dict.keys():
+        raise TypeError("Missing required parameters: 'operation'")
+    operation = inp_dict.get('operation')
+    del inp_dict['operation']
+
+    inputs = [v for k,v in inp_dict.items()]
+    return merge_logical_list(inputs, operation)
+
+
+def merge_logical_list(inp_list: List, operation: Operations):
+    if operation == Operations.OR:
+        return any(inp_list)
+    if operation == Operations.AND:
+        return all(inp_list)
+    raise ValueError("Invalid logical operation")
+
+
+def partial_dict_update(orig_data, new_data):
+    """
+    This function is used when some key value pairs in orig_data should
+    be updated from new_data.
+    """
+    for k, v in new_data.items():
+        orig_data[k] = v
+
+    return orig_data

kube_watch/modules/logic/scheduler.py

@@ -1,74 +1,74 @@
-from datetime import datetime
-from enum import Enum
-
-from prefect import get_run_logger
-logger = get_run_logger()
-
-class IntervalType(Enum):
-    MINUTES = 'minutes'
-    HOURLY = 'hourly'
-    DAILY = 'daily'
-    WEEKLY = 'weekly'
-    MONTHLY = 'monthly'
-    QUARTERLY = 'quarterly'
-    SEMIANNUAL = 'semiannual'
-    YEARLY = 'yearly'
-
-
-def should_run_task(interval_type, interval_value=None, interval_buffer=10, specific_day=None):
-    """
-    The function `should_run_task` determines whether a task should run based on the specified interval
-    type and values.
-    
-    :param interval_type: The `interval_type` parameter specifies the type of interval at which a task
-    should run. It can take on the following values:
-    :param interval_value: The `interval_value` parameter represents the specific value associated with
-    the interval type. For example, if the interval type is `DAILY`, the `interval_value` would
-    represent the specific hour at which the task should run daily. Similarly, for `WEEKLY`, it would
-    represent the specific day
-    :param interval_buffer: The `interval_buffer` parameter in the `should_run_task` function is a
-    default value set to 20 minutes. This provides an acceptable range for a task to get executed. Suitable
-    for Daily, Weekly, Monthly, etc. schedules.
-    :param specific_day: The `specific_day` parameter represents the day of the week when the task
-    should run in the case of a weekly interval. The values for `specific_day` are as follows:
-    :return: The function `should_run_task` takes in various parameters related to different interval
-    types (such as minutes, hourly, daily, weekly, monthly, quarterly, semiannual, yearly) and checks if
-    the current datetime matches the specified interval criteria.
-    """
-
-    now = datetime.now()
-    # Match the interval type
-    if interval_type == IntervalType.MINUTES.value:
-        # Runs every 'interval_value' minutes
-        return now.minute % interval_value == 0
-    
-    if interval_type == IntervalType.HOURLY.value:
-        # Runs every 'interval_value' hours on the hour
-        return now.hour % interval_value == 0 and now.minute < interval_buffer
-    
-    if interval_type == IntervalType.DAILY.value:
-        # Runs once a day at 'interval_value' hour
-        return now.hour == interval_value and now.minute < interval_buffer
-
-    if interval_type == IntervalType.WEEKLY.value:
-        # Runs once a week on 'specific_day' (0=Monday, 6=Sunday)
-        return now.weekday() == specific_day and now.hour == 0 and now.minute < interval_buffer
-    
-    if interval_type == IntervalType.MONTHLY.value:
-        # Runs on the 'interval_value' day of each month
-        return now.day == interval_value and now.hour == 0 and now.minute < interval_buffer
-    
-    if interval_type == IntervalType.QUARTERLY.value:
-        # Runs on the first day of each quarter
-        return now.month % 3 == 1 and now.day == 1 and now.hour == 0 and now.minute < interval_buffer
-    
-    if interval_type == IntervalType.SEMIANNUAL.value:
-        # Runs on the first day of the 1st and 7th month
-        return (now.month == 1 or now.month == 7) and now.day == 1 and now.hour == 0 and now.minute < interval_buffer
-    
-    if interval_type == IntervalType.YEARLY.value:
-        # Runs on the first day of the year
-        return now.month == 1 and now.day == 1 and now.hour == 0 and now.minute < interval_buffer
-
-    return False
-
+from datetime import datetime
+from enum import Enum
+
+from prefect import get_run_logger
+logger = get_run_logger()
+
+class IntervalType(Enum):
+    MINUTES = 'minutes'
+    HOURLY = 'hourly'
+    DAILY = 'daily'
+    WEEKLY = 'weekly'
+    MONTHLY = 'monthly'
+    QUARTERLY = 'quarterly'
+    SEMIANNUAL = 'semiannual'
+    YEARLY = 'yearly'
+
+
+def should_run_task(interval_type, interval_value=None, interval_buffer=10, specific_day=None):
+    """
+    The function `should_run_task` determines whether a task should run based on the specified interval
+    type and values.
+    
+    :param interval_type: The `interval_type` parameter specifies the type of interval at which a task
+    should run. It can take on the following values:
+    :param interval_value: The `interval_value` parameter represents the specific value associated with
+    the interval type. For example, if the interval type is `DAILY`, the `interval_value` would
+    represent the specific hour at which the task should run daily. Similarly, for `WEEKLY`, it would
+    represent the specific day
+    :param interval_buffer: The `interval_buffer` parameter in the `should_run_task` function is a
+    default value set to 20 minutes. This provides an acceptable range for a task to get executed. Suitable
+    for Daily, Weekly, Monthly, etc. schedules.
+    :param specific_day: The `specific_day` parameter represents the day of the week when the task
+    should run in the case of a weekly interval. The values for `specific_day` are as follows:
+    :return: The function `should_run_task` takes in various parameters related to different interval
+    types (such as minutes, hourly, daily, weekly, monthly, quarterly, semiannual, yearly) and checks if
+    the current datetime matches the specified interval criteria.
+    """
+
+    now = datetime.now()
+    # Match the interval type
+    if interval_type == IntervalType.MINUTES.value:
+        # Runs every 'interval_value' minutes
+        return now.minute % interval_value == 0
+    
+    if interval_type == IntervalType.HOURLY.value:
+        # Runs every 'interval_value' hours on the hour
+        return now.hour % interval_value == 0 and now.minute < interval_buffer
+    
+    if interval_type == IntervalType.DAILY.value:
+        # Runs once a day at 'interval_value' hour
+        return now.hour == interval_value and now.minute < interval_buffer
+
+    if interval_type == IntervalType.WEEKLY.value:
+        # Runs once a week on 'specific_day' (0=Monday, 6=Sunday)
+        return now.weekday() == specific_day and now.hour == 0 and now.minute < interval_buffer
+    
+    if interval_type == IntervalType.MONTHLY.value:
+        # Runs on the 'interval_value' day of each month
+        return now.day == interval_value and now.hour == 0 and now.minute < interval_buffer
+    
+    if interval_type == IntervalType.QUARTERLY.value:
+        # Runs on the first day of each quarter
+        return now.month % 3 == 1 and now.day == 1 and now.hour == 0 and now.minute < interval_buffer
+    
+    if interval_type == IntervalType.SEMIANNUAL.value:
+        # Runs on the first day of the 1st and 7th month
+        return (now.month == 1 or now.month == 7) and now.day == 1 and now.hour == 0 and now.minute < interval_buffer
+    
+    if interval_type == IntervalType.YEARLY.value:
+        # Runs on the first day of the year
+        return now.month == 1 and now.day == 1 and now.hour == 0 and now.minute < interval_buffer
+
+    return False
+