supertable 0.1.0__py3-none-any.whl

supertable/config/defaults.py

@@ -0,0 +1,81 @@
+import os
+import logging
+import sys
+
+from dotenv import load_dotenv
+from dataclasses import dataclass
+
+import colorlog
+
+# Configure colors for all levels
+handler = colorlog.StreamHandler()
+handler.setFormatter(colorlog.ColoredFormatter(
+    '%(log_color)s%(asctime)s - %(levelname)-8s - %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S',
+    log_colors={
+        'DEBUG': 'cyan',
+        'INFO': 'green',
+        'WARNING': 'yellow',
+        'ERROR': 'red',
+        'CRITICAL': 'red,bg_white',  # White background with red text
+    },
+    secondary_log_colors={},
+    style='%'
+))
+
+logging.basicConfig(level=logging.INFO, handlers=[handler])
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Default:
+    MAX_MEMORY_CHUNK_SIZE: int = 16 * 1024 * 1024
+    DEFAULT_TIMEOUT_SEC: int = 10
+    DEFAULT_LOCK_DURATION_SEC: int = 60
+    LOG_LEVEL: str = "INFO"  # Replaced IS_DEBUG with LOG_LEVEL
+    IS_SHOW_TIMING: bool = True
+    STORAGE_TYPE: str = "LOCAL"
+
+    def update_default(self, **kwargs):
+        """
+        Updates fields of this Default instance in-place
+        with any matching keys in kwargs.
+        """
+        for key, value in kwargs.items():
+            if hasattr(self, key):
+                setattr(self, key, value)
+                if key == "LOG_LEVEL":
+                    self._update_log_level()
+
+    def _update_log_level(self):
+        """Update the logging level based on the current setting"""
+        logging.getLogger().setLevel(self.LOG_LEVEL)
+        logger.info(f"Log level changed to {self.LOG_LEVEL}")
+
+
+def load_defaults_from_env(env_file: str = ".env") -> Default:
+    load_dotenv(env_file)
+
+    # Get log level from env, default to INFO
+    log_level = os.getenv("LOG_LEVEL", "INFO").upper()
+
+    # Validate log level
+    valid_levels = ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]
+    if log_level not in valid_levels:
+        log_level = "INFO"
+        logger.warning(f"Invalid LOG_LEVEL in .env. Using default INFO. Valid levels are: {valid_levels}")
+
+    # Set the log level immediately
+    logging.getLogger().setLevel(log_level)
+
+    return Default(
+        MAX_MEMORY_CHUNK_SIZE=int(os.getenv("MAX_MEMORY_CHUNK_SIZE", 16 * 1024 * 1024)),
+        DEFAULT_TIMEOUT_SEC=int(os.getenv("DEFAULT_TIMEOUT_SEC", 10)),
+        DEFAULT_LOCK_DURATION_SEC=int(os.getenv("DEFAULT_LOCK_DURATION_SEC", 60)),
+        LOG_LEVEL=log_level,
+        IS_SHOW_TIMING=(os.getenv("IS_SHOW_TIMING", "True").lower() == "true"),
+        STORAGE_TYPE=os.getenv("STORAGE_TYPE", "LOCAL").upper(),
+    )
+
+
+default = load_defaults_from_env()

supertable/config/homedir.py

@@ -0,0 +1,34 @@
+import os
+import sys
+
+from supertable.config.defaults import default, logger
+
+# If this file is located in a subdirectory, adjust the path logic as needed.
+# Currently appending ".." from __file__ to add the project root directory
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+app_home = os.getenv("SUPERTABLE_HOME", "~/supertable")
+
+def change_to_app_home(home_dir: str) -> None:
+    """
+    Attempts to change the current working directory to `home_dir`.
+    Prints the outcome. Logs (or prints) any error encountered.
+    """
+    expanded_dir = os.path.expanduser(home_dir)
+    try:
+        os.chdir(expanded_dir)
+        logger.debug(f"Changed working directory to {expanded_dir}")
+    except Exception as e:
+        logger.error(f"Failed to change working directory to {expanded_dir}: {e}")
+
+if app_home:
+    change_to_app_home(app_home)
+else:
+    logger.error("SUPERTABLE_HOME environment variable is not set")
+
+
+logger.info(f"Current working directory: {os.getcwd()}")
+
+
+def get_app_home():
+    return app_home

supertable/data_reader.py

@@ -0,0 +1,213 @@
+from enum import Enum
+
+import duckdb
+import pandas as pd
+
+from supertable.config.defaults import logger
+from supertable.utils.timer import Timer
+from supertable.super_table import SuperTable
+from supertable.query_plan_manager import QueryPlanManager
+from supertable.utils.sql_parser import SQLParser
+from supertable.utils.helper import dict_keys_to_lowercase
+from supertable.plan_extender import PlanStats, extend_execution_plan
+from supertable.rbac.access_control import restrict_read_access
+
+class Status(Enum):
+    OK = "ok"
+    ERROR = "error"
+
+
+class DataReader:
+    def __init__(self, super_name, organization, query):
+        self.super_table = SuperTable(super_name=super_name, organization=organization)
+        self.parser = SQLParser(query)
+        self.parser.parse_sql()
+        self.timer = None
+        self.plan_stats = None
+        self.query_plan_manager = None
+
+    def filter_snapshots(self, super_table_data, super_table_meta):
+        snapshots = super_table_data.get("snapshots")
+        file_count = super_table_meta.get("file_count", 0)
+        total_rows = super_table_meta.get("total_rows", 0)
+        total_file_size = super_table_meta.get("total_file_size", 0)
+        self.plan_stats.add_stat({"TABLE_FILES": file_count})
+        self.plan_stats.add_stat({"TABLE_SIZE": total_file_size})
+        self.plan_stats.add_stat({"TABLE_ROWS": total_rows})
+
+        if self.super_table.super_name.lower() == self.parser.original_table.lower():
+            filtered_snapshots = [
+                s for s in snapshots
+                if not (s["table_name"].startswith("__") and s["table_name"].endswith("__"))
+            ]
+            return filtered_snapshots
+        else:
+            filtered_snapshots = [
+                entry
+                for entry in snapshots
+                if entry["table_name"].lower() == self.parser.original_table.lower()
+            ]
+
+            return filtered_snapshots
+
+    timer = Timer()
+    @timer
+    def execute(self, user_hash: str, with_scan: bool=False):
+        status = Status.ERROR
+        message = None
+        self.timer = Timer()
+        self.plan_stats = PlanStats()
+
+        try:
+            super_table_data, super_table_path, super_table_meta = self.super_table.get_super_table_and_path_with_shared_lock()
+
+            self.timer.capture_and_reset_timing(event="META")
+
+            self.query_plan_manager = QueryPlanManager(super_name=self.super_table.super_name,
+                                                       organization=self.super_table.organization,
+                                                       current_meta_path=super_table_path,
+                                                       parser=self.parser)
+
+            snapshots = self.filter_snapshots(super_table_data=super_table_data,
+                                              super_table_meta=super_table_meta)
+            logger.debug(f"Filtered snapshots: {snapshots}")
+
+            parquet_files, schema = self.process_snapshots(snapshots=snapshots,
+                                                           with_scan=with_scan)
+            logger.debug(f"Parquet Files: {parquet_files}")
+
+            missing_columns = (
+                set([column.lower() for column in self.parser.columns_list])
+                - set("*")
+                - schema
+            )
+            logger.debug(f"Mising Columns: {missing_columns}")
+
+            if len(snapshots) == 0 or missing_columns or not parquet_files:
+                message = (
+                    f"Missing column(s): {', '.join ( missing_columns )}"
+                    if missing_columns
+                    else "No parquet files found"
+                )
+                logger.warning(f"Filter Result: {message}")
+                return pd.DataFrame(), status, message
+
+            restrict_read_access(super_name=self.super_table.super_name,
+                                 organization=self.super_table.organization,
+                                 user_hash=user_hash,
+                                 table_name=self.parser.reflection_table,
+                                 table_schema=schema,
+                                 parsed_columns=self.parser.columns_list,
+                                 parser=self.parser)
+
+            self.timer.capture_and_reset_timing(event="FILTERING")
+
+            result = self.execute_with_duckdb(parquet_files=parquet_files,
+                                              query_manager=self.query_plan_manager)
+
+            status = Status.OK
+        except Exception as e:
+            message = str(e)
+            logger.error(f"Exception: {e}")
+            result = pd.DataFrame()
+        self.timer.capture_and_reset_timing(event="EXECUTING_QUERY")
+
+        try:
+            extend_execution_plan(super_table=self.super_table,
+                                  user_hash=user_hash,
+                                  query_plan_manager=self.query_plan_manager,
+                                  timing=self.timer.timings,
+                                  status=status.value,
+                                  message=message,
+                                  result_shape=result.shape,
+                                  plan_stats=self.plan_stats)
+        except Exception as e:
+            logger.error(f"Exception: {e}")
+
+        self.timer.capture_and_reset_timing(event="EXTENDING_PLAN")
+        self.timer.capture_duration(event="TOTAL_EXECUTE")
+        return result, status, message
+
+    def process_snapshots(self, snapshots, with_scan):
+        parquet_files = []
+        reflection_file_size = 0
+        reflection_rows = 0
+
+        schema = set()
+        for snapshot in snapshots:
+            current_snapshot_path = snapshot["path"]
+            current_snapshot_data = self.super_table.read_simple_table_snapshot(
+                current_snapshot_path
+            )
+
+            current_schema = current_snapshot_data.get("schema", {})
+            resources = current_snapshot_data.get("resources", {})
+            schema.update(dict_keys_to_lowercase(current_schema).keys())
+
+            for resource in resources:
+                file_size = resource.get("file_size", 0)
+                file_rows = resource.get("rows", 0)
+
+                if (
+                    with_scan
+                    or self.parser.columns_csv == "*"
+                    or any(
+                        col in dict_keys_to_lowercase(current_schema).keys()
+                        for col in [
+                            column.lower() for column in self.parser.columns_list
+                        ]
+                    )
+                ):
+                    parquet_files.append(resource["file"])
+                    reflection_file_size += file_size
+                    reflection_rows += file_rows
+
+            logger.debug(f"snapshots: {len ( snapshots )}")
+            logger.debug(f"parquet_files: {len ( parquet_files )}")
+            logger.debug(f"schema: {schema}")
+
+        self.plan_stats.add_stat({"REFLECTIONS": len(parquet_files)})
+        self.plan_stats.add_stat({"REFLECTION_SIZE": reflection_file_size})
+        self.plan_stats.add_stat({"REFLECTION_ROWS": reflection_rows})
+
+        return parquet_files, schema
+
+    def execute_with_duckdb(self, parquet_files, query_manager: QueryPlanManager):
+        # Use DuckDB to read and query the parquet files directly
+        con = duckdb.connect()
+
+        con.execute("PRAGMA memory_limit='2GB';")
+        con.execute(f"PRAGMA temp_directory='{query_manager.temp_dir}';")
+        con.execute("PRAGMA enable_profiling='json';")
+        #con.execute("SET profiling_mode = 'standard';")
+        con.execute(f"PRAGMA profile_output = '{query_manager.query_plan_path}';")
+
+        # Read and register parquet files directly with DuckDB
+        parquet_files_str = ", ".join(f"'{file}'" for file in parquet_files)
+        logger.debug(f"parquet files: {len(parquet_files)}")
+
+        self.timer.capture_and_reset_timing("CONNECTING")
+
+        create_table = f"""
+CREATE TABLE {self.parser.reflection_table} 
+AS 
+SELECT {self.parser.columns_csv}
+FROM parquet_scan([{parquet_files_str}], union_by_name=True, HIVE_PARTITIONING=TRUE);
+          """
+
+        logger.debug(f"create_table: {create_table}")
+        con.execute(create_table)
+
+        create_view = f"""
+CREATE VIEW {self.parser.rbac_view}
+AS
+{self.parser.view_definition}
+"""
+        logger.debug(f"create_view: {create_view}")
+        con.execute(create_view)
+
+        self.timer.capture_and_reset_timing("CREATING_REFLECTION")
+        logger.debug(f"Executing Query: {self.parser.executing_query}")
+        result = con.execute(query=self.parser.executing_query).fetchdf()
+        logger.debug(f"result.shape: {result.shape}")
+        return result

supertable/data_writer.py

@@ -0,0 +1,100 @@
+import polars
+import re
+
+from polars import DataFrame
+
+from supertable.config.defaults import logger
+from supertable.super_table import SuperTable
+from supertable.simple_table import SimpleTable
+from supertable.utils.timer import Timer
+from supertable.processing import (
+    process_overlapping_files,
+    find_and_lock_overlapping_files,
+)
+from supertable.rbac.access_control import check_write_access
+
+class DataWriter:
+    def __init__(self, super_name: str, organization: str):
+        self.super_table = SuperTable(super_name, organization)
+
+    timer = Timer()
+
+    @timer
+    def write(self, user_hash, simple_name, data, overwrite_columns, compression_level=1):
+
+        logger.debug("Checking for Write Access")
+        check_write_access(super_name=self.super_table.super_name,
+                           organization=self.super_table.organization,
+                           user_hash=user_hash,
+                           table_name=simple_name)
+        logger.debug("Passed Write Access Check")
+
+        # Convert the input dataset from Arrow format to a Polars DataFrame
+        logger.debug("Converting data to DataFrame")
+        dataframe: DataFrame = polars.from_arrow(data)
+        logger.debug("Converted data to DataFrame")
+
+        logger.debug("Validating the dataframe")
+        self.validation(dataframe, simple_name, overwrite_columns)
+        logger.debug("dataframe is valid")
+
+        logger.debug(f"Reading Simple Table Metadata {simple_name}")
+        simple_table = SimpleTable(self.super_table, simple_name)
+        last_simple_table, _ = simple_table.get_simple_table_with_shared_lock()
+        logger.debug(f"last_simple_table: {last_simple_table}")
+
+        # Find files that have overlapping data and lock them to prevent concurrent modifications
+        overlapping_files = find_and_lock_overlapping_files(
+            last_simple_table, dataframe, overwrite_columns, simple_table.locking
+        )
+        logger.debug(f"overlapping_files: {overlapping_files}")
+
+        # Process the overlapping files by filtering, merging, and updating resources
+        inserted, deleted, total_rows, total_columns, new_resources, sunset_files = (
+            process_overlapping_files(
+                dataframe,
+                overlapping_files,
+                overwrite_columns,
+                simple_table.data_dir,
+                compression_level,
+            )
+        )
+
+        new_simple_table_snapshot, new_simple_table_path = simple_table.lock_and_update(
+            new_resources, sunset_files, dataframe
+        )
+
+        self.super_table.update_with_lock(
+            simple_name, new_simple_table_path, new_simple_table_snapshot
+        )
+
+        simple_table.locking.release_lock()
+        # self.super_table.locking.release_lock()
+
+        return total_columns, total_rows, inserted, deleted
+
+    def validation(
+        self, dataframe: DataFrame, simple_name: str, overwrite_columns: list
+    ):
+        if len(simple_name) == 0 or len(simple_name) > 128:
+            raise ValueError("SimpleTable name can't be empty or longer than 128")
+
+        if simple_name == self.super_table.super_name:
+            raise ValueError("SimpleTable name can't match with SuperTable name")
+
+        # Regular expression pattern for a valid table name
+        pattern = r"^[A-Za-z_][A-Za-z0-9_]*$"
+        if not re.match(pattern, simple_name):
+            raise ValueError(
+                f"Invalid table name: '{simple_name}'. Table names must start with a letter or underscore and contain only alphanumeric characters and underscores."
+            )
+
+        # Validate the overwrite columns
+        if overwrite_columns and not all(
+            col in dataframe.columns for col in overwrite_columns
+        ):
+            raise ValueError("Some overwrite columns are not present in the dataset")
+
+        # Ensure overwrite_columns is a list
+        if isinstance(overwrite_columns, str):
+            raise ValueError("overwrite columns must be list")

supertable/history_cleaner.py

@@ -0,0 +1,99 @@
+import os
+from supertable.config.defaults import logger
+from supertable.super_table import SuperTable
+from supertable.rbac.access_control import check_write_access
+
+class HistoryCleaner:
+    def __init__(self, super_name: str, organization: str):
+        self.super_table = SuperTable(super_name=super_name, organization=organization)
+        # Grab the storage object from the super_table
+        self.storage = self.super_table.storage
+
+    def clean(self, user_hash):
+        # Acquire the shared lock and read super table meta
+        super_table_data, super_table_path, super_table_meta = (
+            self.super_table.get_super_table_and_path_with_shared_lock()
+        )
+        last_updated_ms = super_table_data["last_updated_ms"]
+
+        check_write_access(super_name=self.super_table.super_name,
+                           organization=self.super_table.organization,
+                           user_hash=user_hash,
+                           table_name=self.super_table.super_name)
+
+        # Collect all files in the super_table.super_dir
+        super_files = self.collect_files(self.super_table.super_dir)
+        # Remove the current super_table pointer from the list
+        if super_table_path in super_files:
+            super_files.remove(super_table_path)
+
+        files_to_delete = self.get_files_to_delete(super_files, last_updated_ms)
+        self.delete_files(files_to_delete)
+
+        logger.debug(
+            f"{len(files_to_delete)} files cleaned for table: {self.super_table.super_name}"
+        )
+
+        # Iterate over each snapshot in the super table
+        for snapshot in super_table_data["snapshots"]:
+            # Note: If your snapshot dict uses "table_name" instead of "simple_name", adjust here:
+            simple_table_name = snapshot["table_name"]    # or snapshot["simple_name"] if appropriate
+            simple_table_path = snapshot["path"]
+
+            # Read the simple table's own snapshot JSON
+            simple_table_data = self.storage.read_json(simple_table_path)
+
+            location = simple_table_data["location"]
+            active_files = [entry["file"] for entry in simple_table_data.get("resources", [])]
+
+            designated_files = set(self.collect_files(location)) - set(active_files)
+            # Remove the current snapshot path from the designated_files
+            if simple_table_path in designated_files:
+                designated_files.remove(simple_table_path)
+
+            # Convert back to list
+            designated_files = list(designated_files)
+
+            files_to_delete = self.get_files_to_delete(designated_files, last_updated_ms)
+            self.delete_files(files_to_delete)
+
+            logger.debug(f"{len(files_to_delete)} files cleaned for table: {simple_table_name}")
+
+    def collect_files(self, location):
+        """
+        Collect parquet & JSON files under 'location'.
+        This uses the storage interface's list_files instead of glob.
+        """
+        # Example usage for a typical directory layout:
+        parquet_files = self.storage.list_files(os.path.join(location, "data"), "*.parquet")
+        json_files = self.storage.list_files(os.path.join(location, "snapshots"), "*.json")
+        super_json_files = self.storage.list_files(location, "*.json")
+        return parquet_files + json_files + super_json_files
+
+    def get_files_to_delete(self, designated_files, last_updated_ms):
+        """
+        Compare the numeric timestamp in the file name (e.g. 1678900000_*.json)
+        against `last_updated_ms`.
+        """
+        files_to_delete = []
+        for file in designated_files:
+            # Example: "1678900000_fileinfo.json" => "1678900000"
+            filename = os.path.basename(file)
+            timestamp_str = filename.split("_")[0]
+            try:
+                timestamp_val = int(timestamp_str)
+                if timestamp_val <= last_updated_ms:
+                    files_to_delete.append(file)
+            except ValueError:
+                # If the file doesn't match the pattern, skip or handle differently
+                pass
+
+        return files_to_delete
+
+    def delete_files(self, files_to_delete):
+        """
+        Deletes files using the storage interface's `delete` method.
+        """
+        for file in files_to_delete:
+            self.storage.delete(file)
+            logger.debug(f"Deleted file: {file}")

supertable/locking/__init__.py

@@ -0,0 +1,3 @@
+from supertable.locking.locking import Locking
+
+__all__ = ["Locking"]

supertable/locking/file_lock.py

@@ -0,0 +1,104 @@
+import json
+import os
+import secrets
+import time
+import fcntl
+from supertable.config.defaults import default, logger
+
+class FileLocking:
+    def __init__(self, identity, working_dir, lock_file_name=".lock.json", check_interval=0.1):
+        self.identity = identity
+        self.lock_id = secrets.token_hex(8)
+        self.lock_file_dir = working_dir if working_dir is not None else identity
+        self.lock_file_path = os.path.join(self.lock_file_dir, lock_file_name)
+        self.check_interval = check_interval
+        logger.debug(f"lock_file_dir: {self.lock_file_dir}")
+        logger.debug(f"lock_file_path: {self.lock_file_path}")
+        self.init_lock_file()
+
+    def init_lock_file(self):
+        os.makedirs(self.lock_file_dir, exist_ok=True)
+        if not os.path.exists(self.lock_file_path):
+            with open(self.lock_file_path, "w") as lock_file:
+                json.dump([], lock_file)
+
+    def read_lock_file(self, lock_file):
+        lock_file.seek(0)
+        try:
+            return json.load(lock_file)
+        except json.JSONDecodeError as e:
+            logger.error(f"Error reading lock file: {e}")
+            return []
+
+    def write_lock_file(self, lock_data, lock_file):
+        lock_file.seek(0)
+        lock_file.truncate()
+        json.dump(lock_data, lock_file)
+        lock_file.flush()
+        os.fsync(lock_file.fileno())
+
+    def remove_expired_locks(self, lock_data):
+        current_time = int(time.time())
+        return [lock for lock in lock_data if lock["exp"] > current_time]
+
+    def remove_own_locks(self, lock_data):
+        return [lock for lock in lock_data if lock["pid"] != self.lock_id]
+
+    def self_lock(self, timeout_seconds=default.DEFAULT_TIMEOUT_SEC, lock_duration_seconds=default.DEFAULT_LOCK_DURATION_SEC):
+        resources = [self.identity]
+        return self.lock_resources(resources, timeout_seconds, lock_duration_seconds)
+
+    def lock_resources(self, resources, timeout_seconds=default.DEFAULT_TIMEOUT_SEC, lock_duration_seconds=default.DEFAULT_LOCK_DURATION_SEC):
+        start_time = time.time()
+        expiration_time = int(time.time() + lock_duration_seconds)
+        sleep_time = 0
+
+        while time.time() - start_time < timeout_seconds:
+            try:
+                if sleep_time > 0:
+                    logger.debug(f"Waiting {sleep_time} seconds to acquire the lock for {self.identity}")
+                    time.sleep(sleep_time)
+
+                with open(self.lock_file_path, "r+") as lock_file:
+                    fcntl.flock(lock_file, fcntl.LOCK_EX)
+                    try:
+                        lock_data = self.read_lock_file(lock_file)
+                        lock_data = self.remove_expired_locks(lock_data)
+                        lock_check = self.remove_own_locks(lock_data)
+                        logger.debug(f"Lock data: {lock_data}")
+                        logger.debug(f"Lock check: {lock_check}")
+
+                        if any(resource in lock["res"] for lock in lock_check for resource in resources):
+                            logger.debug(f"{self.identity}: lock can't be acquired for resources {resources}")
+                            sleep_time = self.check_interval
+                            continue
+
+                        lock_entry = {"pid": self.lock_id, "exp": expiration_time, "res": resources}
+                        lock_data.append(lock_entry)
+                        self.write_lock_file(lock_data, lock_file)
+                        logger.debug(f"Identity: {self.identity}, lock acquired: {self.lock_id}")
+                        return True
+                    finally:
+                        fcntl.flock(lock_file, fcntl.LOCK_UN)
+            except Exception as e:
+                logger.error(f"Error during lock acquisition: {e}")
+                time.sleep(self.check_interval)
+        return False
+
+    def release_lock(self, resources=None):
+        with open(self.lock_file_path, "r+") as lock_file:
+            fcntl.flock(lock_file, fcntl.LOCK_EX)
+            try:
+                lock_data = self.read_lock_file(lock_file)
+                lock_data = [lock for lock in lock_data if lock["pid"] != self.lock_id]
+                self.write_lock_file(lock_data, lock_file)
+            finally:
+                fcntl.flock(lock_file, fcntl.LOCK_UN)
+
+    def __enter__(self):
+        if not self.self_lock():
+            raise Exception(f"Unable to acquire file lock for {self.identity}")
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.release_lock()