adss 0.1__py3-none-any.whl → 1.1__py3-none-any.whl

adss/models/query.py

@@ -0,0 +1,134 @@
+"""
+Query-related data models for the Astronomy TAP Client.
+"""
+from dataclasses import dataclass
+from typing import Dict, Optional, Any, List
+from datetime import datetime
+import pandas as pd
+
+from ..utils import parse_datetime
+
+
+@dataclass
+class Query:
+    """
+    Represents a database query and its metadata.
+    """
+    id: str
+    query_text: str
+    status: str  # 'PENDING', 'QUEUED', 'RUNNING', 'COMPLETED', 'ERROR'
+    created_at: datetime
+    mode: str = 'adql'  # 'adql' or 'sql'
+    user_id: Optional[str] = None
+    completed_at: Optional[datetime] = None
+    result_url: Optional[str] = None
+    error: Optional[str] = None
+    execution_time_ms: Optional[int] = None
+    row_count: Optional[int] = None
+    position_in_queue: Optional[int] = None
+    expires_at: Optional[datetime] = None
+    query_metadata: Optional[Dict[str, Any]] = None
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> 'Query':
+        """Create a Query object from a dictionary."""
+        query_id = data.get('id')
+        query_text = data.get('query_text')
+        status = data.get('status')
+        mode = data.get('mode', 'adql')
+        user_id = data.get('user_id')
+        
+        created_at = parse_datetime(data.get('created_at'))
+        completed_at = parse_datetime(data.get('completed_at'))
+        expires_at = parse_datetime(data.get('expires_at'))
+        
+        result_url = data.get('result_url')
+        error = data.get('error')
+        execution_time_ms = data.get('execution_time_ms')
+        row_count = data.get('row_count')
+        position_in_queue = data.get('position_in_queue')
+        query_metadata = data.get('query_metadata')
+        
+        return cls(
+            id=query_id,
+            query_text=query_text,
+            status=status,
+            mode=mode,
+            user_id=user_id,
+            created_at=created_at,
+            completed_at=completed_at,
+            result_url=result_url,
+            error=error,
+            execution_time_ms=execution_time_ms,
+            row_count=row_count,
+            position_in_queue=position_in_queue,
+            expires_at=expires_at,
+            query_metadata=query_metadata
+        )
+    
+    @property
+    def is_complete(self) -> bool:
+        """Check if the query has completed (successfully or with error)."""
+        return self.status in ['COMPLETED', 'ERROR']
+    
+    @property
+    def is_running(self) -> bool:
+        """Check if the query is currently running."""
+        return self.status == 'RUNNING'
+    
+    @property
+    def is_queued(self) -> bool:
+        """Check if the query is queued."""
+        return self.status == 'QUEUED'
+    
+    @property
+    def is_successful(self) -> bool:
+        """Check if the query completed successfully."""
+        return self.status == 'COMPLETED'
+    
+    @property
+    def is_failed(self) -> bool:
+        """Check if the query failed."""
+        return self.status == 'ERROR'
+
+
+@dataclass
+class QueryResult:
+    """
+    Represents the result of a query, including the data and metadata.
+    """
+    query: Query
+    data: pd.DataFrame
+    execution_time_ms: Optional[int] = None
+    row_count: Optional[int] = None
+    column_count: Optional[int] = None
+    
+    def to_csv(self, path: str, **kwargs) -> None:
+        """Save the query result to a CSV file."""
+        self.data.to_csv(path, **kwargs)
+    
+    def to_parquet(self, path: str, **kwargs) -> None:
+        """Save the query result to a Parquet file."""
+        self.data.to_parquet(path, **kwargs)
+    
+    def to_json(self, path: str = None, **kwargs) -> Optional[str]:
+        """
+        Convert the query result to JSON.
+        If path is provided, saves to file, otherwise returns a JSON string.
+        """
+        if path:
+            self.data.to_json(path, **kwargs)
+            return None
+        return self.data.to_json(**kwargs)
+    
+    def head(self, n: int = 5) -> pd.DataFrame:
+        """Return the first n rows of the result."""
+        return self.data.head(n)
+    
+    def tail(self, n: int = 5) -> pd.DataFrame:
+        """Return the last n rows of the result."""
+        return self.data.tail(n)
+    
+    def describe(self) -> pd.DataFrame:
+        """Return summary statistics of the result."""
+        return self.data.describe()

adss/models/user.py

@@ -0,0 +1,123 @@
+"""
+User-related data models for the Astronomy TAP Client.
+"""
+from dataclasses import dataclass, field
+from typing import List, Dict, Optional, Any
+from datetime import datetime
+
+from ..utils import parse_datetime
+
+
+@dataclass
+class SchemaPermission:
+    """Schema-level permission."""
+    schema_name: str
+    permission: str  # 'read', 'write', or 'all'
+
+
+@dataclass
+class TablePermission:
+    """Table-level permission."""
+    schema_name: str
+    table_name: str
+    permission: str  # 'read', 'write', or 'all'
+
+
+@dataclass
+class RolePermissions:
+    """Permissions associated with a role."""
+    schema_permissions: List[SchemaPermission] = field(default_factory=list)
+    table_permissions: List[TablePermission] = field(default_factory=list)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> 'RolePermissions':
+        """Create a RolePermissions object from a dictionary."""
+        schema_perms = [
+            SchemaPermission(**p) 
+            for p in data.get('schema_permissions', [])
+        ]
+        
+        table_perms = [
+            TablePermission(**p) 
+            for p in data.get('table_permissions', [])
+        ]
+        
+        return cls(
+            schema_permissions=schema_perms,
+            table_permissions=table_perms
+        )
+
+
+@dataclass
+class Role:
+    """User role with associated permissions."""
+    id: int
+    name: str
+    description: Optional[str] = None
+    permissions: Optional[RolePermissions] = None
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> 'Role':
+        """Create a Role object from a dictionary."""
+        role_id = data.get('id')
+        name = data.get('name')
+        description = data.get('description')
+        
+        permissions_data = data.get('permissions')
+        permissions = None
+        if permissions_data:
+            permissions = RolePermissions.from_dict(permissions_data)
+        
+        return cls(
+            id=role_id,
+            name=name,
+            description=description,
+            permissions=permissions
+        )
+
+
+@dataclass
+class User:
+    """User model with authentication and role information."""
+    id: str
+    username: str
+    email: str
+    full_name: Optional[str] = None
+    is_active: bool = True
+    is_staff: bool = False
+    is_superuser: bool = False
+    created_at: Optional[datetime] = None
+    last_login: Optional[datetime] = None
+    roles: List[Role] = field(default_factory=list)
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> 'User':
+        """Create a User object from a dictionary."""
+        user_id = data.get('id')
+        username = data.get('username')
+        email = data.get('email')
+        full_name = data.get('full_name')
+        is_active = data.get('is_active', True)
+        is_staff = data.get('is_staff', False)
+        is_superuser = data.get('is_superuser', False)
+        
+        created_at = parse_datetime(data.get('created_at'))
+        last_login = parse_datetime(data.get('last_login'))
+        
+        roles = [
+            Role.from_dict(role_data) 
+            for role_data in data.get('roles', [])
+        ]
+        
+        return cls(
+            id=user_id,
+            username=username,
+            email=email,
+            full_name=full_name,
+            is_active=is_active,
+            is_staff=is_staff,
+            is_superuser=is_superuser,
+            created_at=created_at,
+            last_login=last_login,
+            roles=roles
+        )

adss/table.py

@@ -0,0 +1,295 @@
+from adss.executors.sync_query import execute_sync
+from adss.executors.async_query import execute_async
+from adss.utils import format_table
+
+import re
+
+class Table:
+    def __init__(self, name, columns):
+        self.name = name
+        self.columns = columns
+        self.selected_columns = []
+        self.constrains = []
+        
+    def __repr__(self):
+        return f"Table(name={self.name}, columns={len(self.columns)})"
+    
+    def __str__(self):
+        return f"Table: {self.name} ({len(self.columns)} columns)"
+    
+    def check_column(self, column):
+        return column in self.columns
+    
+    def format_columns(self, columns):
+        # Use provided columns list rather than an undefined variable
+        return ','.join(columns)
+    
+    def set_columns(self, columns):
+        if not isinstance(columns, list):
+            columns = [columns]
+        for column in columns:
+            if not self.check_column(column):
+                raise ValueError(f"Column {column} not in table {self.name}, options are {self.columns}")
+        self.selected_columns = columns
+    
+    def set_constrains(self, constrains):
+        self.constrains = constrains
+    
+    def cone_search(self, ra, dec, radius_arcsec, columns=None, method = 'sync'):
+        if radius_arcsec < 0:
+            raise ValueError("Radius must be positive")
+        if radius_arcsec > 60:
+            raise ValueError("Radius must be less than 60 arcsecs")
+        
+        if columns:
+            columns_str = self.format_columns(columns)
+        elif self.selected_columns:
+            columns_str = self.format_columns(self.selected_columns)
+        else:
+            columns_str = "*"  # Select all columns
+
+        constraints_str = ""
+        if self.constrains:
+            constraints_str = " AND (" + self.constrains + ")"
+            
+        query = f"""SELECT {columns_str} FROM {self.name}
+WHERE 1 = CONTAINS( 
+    POINT('ICRS', ra, dec), 
+    CIRCLE('ICRS', {ra}, {dec}, {radius_arcsec}/3600.0) 
+){constraints_str}
+"""
+                
+        print(query)
+        if method == 'sync':
+            return execute_sync(query)
+        else:
+            return execute_async(query)
+    
+    def cone_cross_match(
+            self, 
+            other_table, 
+            match_arcsec,
+            ra, 
+            dec,
+            radius_arcsec,
+            columns=None, 
+            other_columns=None,
+            other_suffix=None,
+            method='sync'
+        ):
+        """
+        Perform a cone search on the current table (t1) and then cross-match with another table (t2)
+        using a matching radius (match_arcsec).
+        
+        The query first restricts table t1 to a cone centered at (ra, dec) with a radius of radius_arcsec.
+        Then, for each object in t1, it finds matching objects in table t2 that lie within match_arcsec 
+        of the t1 object's coordinates.
+        
+        Additionally:
+          - If a non-empty `other_suffix` is provided, each selected column from t2 will be aliased with that suffix.
+          - The constraints for each table are processed so that the columns in the conditions are properly qualified with t1 or t2.
+        
+        Parameters:
+            other_table (Table): The table to match against (t2).
+            match_arcsec (float): The cross-match tolerance radius (in arcseconds) between t1 and t2.
+            ra (float): Right Ascension for the cone center (t1).
+            dec (float): Declination for the cone center (t1).
+            radius_arcsec (float): The cone search radius (in arcseconds) for filtering t1.
+            columns (list or None): Columns to select from the current table (t1).
+            other_columns (list or None): Columns to select from the other table (t2).
+            other_suffix (str or None): Optional suffix to append to each t2 column alias.
+            method (str): Use 'sync' for synchronous execution or 'async' for asynchronous.
+        
+        Returns:
+            The result of the query execution via execute_sync or execute_async.
+        """
+        
+        # Helper function to qualify constraint column names with the proper alias.
+        # It looks for each column name as a whole word (not already preceded by an alias) and prefixes it.
+        def apply_alias_to_constraint(constraint, alias, columns_list):
+            for col in columns_list:
+                # (?<![\w\.]) ensures that we do not match if the column is already prefixed (like t1.ra)
+                pattern = r'(?<![\w\.])\b' + re.escape(col) + r'\b'
+                constraint = re.sub(pattern, f"{alias}.{col}", constraint)
+            return constraint
+        
+        # Validate match_arcsec
+        if match_arcsec <= 0:
+            raise ValueError("Match radius must be positive")
+        if match_arcsec > 3:
+            print("Match radius may be too large; consider a value less than 3 arcsecs")
+        
+        # Determine columns for t1
+        if columns:
+            t1_columns_list = columns if isinstance(columns, list) else [columns]
+            t1_columns = ', '.join(f"t1.{col}" for col in t1_columns_list)
+        elif self.selected_columns:
+            t1_columns = ', '.join(f"t1.{col}" for col in self.selected_columns)
+        else:
+            t1_columns = "t1.*"
+        
+        # Determine columns for t2, adding suffix if provided
+        if other_columns:
+            t2_columns_list = other_columns if isinstance(other_columns, list) else [other_columns]
+            if other_suffix:
+                t2_columns = ', '.join(f"t2.{col} AS {col}{other_suffix}" for col in t2_columns_list)
+            else:
+                t2_columns = ', '.join(f"t2.{col}" for col in t2_columns_list)
+        elif other_table.selected_columns:
+            if other_suffix:
+                t2_columns = ', '.join(f"t2.{col} AS {col}{other_suffix}" for col in other_table.selected_columns)
+            else:
+                t2_columns = ', '.join(f"t2.{col}" for col in other_table.selected_columns)
+        else:
+            t2_columns = "t2.*"
+        
+        # Process constraints for t1: apply alias "t1" to each column mentioned in the constraint.
+        constraints_t1 = ""
+        if self.constrains:
+            if isinstance(self.constrains, str):
+                processed_constraint = apply_alias_to_constraint(self.constrains, "t1", self.columns)
+                constraints_t1 = " AND (" + processed_constraint + ")"
+            elif isinstance(self.constrains, list):
+                processed_constraints = []
+                for c in self.constrains:
+                    processed_constraints.append(apply_alias_to_constraint(c, "t1", self.columns))
+                constraints_t1 = " AND (" + " AND ".join(processed_constraints) + ")"
+
+        # Process constraints for t2: apply alias "t2" to each column mentioned in the constraint.
+        constraints_t2 = ""
+        if other_table.constrains:
+            if isinstance(other_table.constrains, str):
+                processed_constraint = apply_alias_to_constraint(other_table.constrains, "t2", other_table.columns)
+                constraints_t2 = " AND (" + processed_constraint + ")"
+            elif isinstance(other_table.constrains, list):
+                processed_constraints = []
+                for c in other_table.constrains:
+                    processed_constraints.append(apply_alias_to_constraint(c, "t2", other_table.columns))
+                constraints_t2 = " AND (" + " AND ".join(processed_constraints) + ")"
+        # Build the query:
+        # 1. The first CONTAINS clause performs the cross-match between t1 and t2 with match_arcsec tolerance.
+        # 2. The second CONTAINS clause restricts t1 objects to the cone centered at (ra, dec) with radius radius_arcsec.
+        query = f"""SELECT {t1_columns}, {t2_columns}
+FROM {self.name} AS t1, {other_table.name} AS t2
+WHERE 1 = CONTAINS(
+        POINT('ICRS', t2.ra, t2.dec),
+        CIRCLE('ICRS', t1.ra, t1.dec, {match_arcsec}/3600.0)
+    )
+AND 1 = CONTAINS(
+        POINT('ICRS', t1.ra, t1.dec),
+        CIRCLE('ICRS', {ra}, {dec}, {radius_arcsec}/3600.0)
+    )
+{constraints_t1}
+{constraints_t2}
+"""
+                
+        print(query)
+        if method == 'async':
+            return execute_async(query)
+        else:
+            return execute_sync(query)
+        
+    def table_cross_match(
+            self, 
+            other_table, 
+            match_arcsec,
+            columns=None, 
+            other_columns=None,
+            other_suffix=None,
+            method='async'
+        ):
+        """
+        Perform a cone search on the current table (t1) and then cross-match with another 
+        table (Dataframe or astropy Table) (t2)
+        using a matching radius (match_arcsec).
+        
+        For each object in t1, it finds matching objects in table t2 that lie within match_arcsec 
+        of the t1 object's coordinates.
+        
+        Additionally:
+          - If a non-empty `other_suffix` is provided, each selected column from t2 will be aliased with that suffix.
+          - The constraints for each table are processed so that the columns in the conditions are properly qualified with t1 or t2.
+        
+        Parameters:
+            other_table (astropy.table.Table): The table to match against (t2).
+            match_arcsec (float): The cross-match tolerance radius (in arcseconds) between t1 and t2.
+            columns (list or None): Columns to select from the current table (t1).
+            other_columns (list or None): Columns to select from the other table (t2).
+            other_suffix (str or None): Optional suffix to append to each t2 column alias.
+            method (str): Use 'sync' for synchronous execution or 'async' for asynchronous.
+        
+        Returns:
+            The result of the query execution via execute_sync or execute_async.
+        """
+        
+        # Helper function to qualify constraint column names with the proper alias.
+        # It looks for each column name as a whole word (not already preceded by an alias) and prefixes it.
+        def apply_alias_to_constraint(constraint, alias, columns_list):
+            for col in columns_list:
+                # (?<![\w\.]) ensures that we do not match if the column is already prefixed (like t1.ra)
+                pattern = r'(?<![\w\.])\b' + re.escape(col) + r'\b'
+                constraint = re.sub(pattern, f"{alias}.{col}", constraint)
+            return constraint
+        
+        # Validate match_arcsec
+        if match_arcsec <= 0:
+            raise ValueError("Match radius must be positive")
+        if match_arcsec > 3:
+            print("Match radius may be too large; consider a value less than 3 arcsecs")
+        
+        # Determine columns for t1
+        if columns:
+            t1_columns_list = columns if isinstance(columns, list) else [columns]
+            t1_columns = ', '.join(f"t1.{col}" for col in t1_columns_list)
+        elif self.selected_columns:
+            t1_columns = ', '.join(f"t1.{col}" for col in self.selected_columns)
+        else:
+            t1_columns = "t1.*"
+        
+        # Determine columns for t2, adding suffix if provided
+        if not other_columns:
+            raise ValueError("Must provide columns for the input table (other_columns param)")
+        
+        if not "ra" in other_columns or not "dec" in other_columns:
+            raise ValueError("Input table must have 'ra' and 'dec' columns")
+        
+        other_table = other_table[other_columns]
+        
+        t2_columns_list = other_columns if isinstance(other_columns, list) else [other_columns]
+        if other_suffix:
+            t2_columns = ', '.join(f"t2.{col} AS {col}{other_suffix}" for col in t2_columns_list)
+        else:
+            t2_columns = ', '.join(f"t2.{col}" for col in t2_columns_list)
+
+        # Process constraints for t1: apply alias "t1" to each column mentioned in the constraint.
+        constraints_t1 = ""
+        if self.constrains:
+            if isinstance(self.constrains, str):
+                processed_constraint = apply_alias_to_constraint(self.constrains, "t1", self.columns)
+                constraints_t1 = " (" + processed_constraint + ")"
+            elif isinstance(self.constrains, list):
+                processed_constraints = []
+                for c in self.constrains:
+                    processed_constraints.append(apply_alias_to_constraint(c, "t1", self.columns))
+                constraints_t1 = " (" + " AND ".join(processed_constraints) + ")"
+
+        if constraints_t1:
+            constraints_t1 = "WHERE " + constraints_t1
+        # Build the query:
+        # 1. The first CONTAINS clause performs the cross-match between t1 and t2 with match_arcsec tolerance.
+        # 2. The second CONTAINS clause restricts t1 objects to the cone centered at (ra, dec) with radius radius_arcsec.
+        query = f"""SELECT {t1_columns}, {t2_columns}
+FROM {self.name} AS t1 JOIN tap_upload.upload AS t2 ON
+1 = CONTAINS(
+        POINT('ICRS', t1.ra, t1.dec),
+        CIRCLE('ICRS', t2.ra, t2.dec, {match_arcsec}/3600.0)
+    )
+{constraints_t1}
+"""
+                
+        print(query)
+        if method == 'async':
+            return execute_async(query, table_upload=other_table)
+        else:
+            raise ValueError("Synchronous execution not supported yet for table cross-match")
+            #return execute_sync(query)

adss/utils/format_table.py

@@ -0,0 +1,115 @@
+import numpy as np
+from astropy.table import Table
+
+def vectorized_string_to_masked_array(column_data):
+    """
+    Fully vectorized conversion of formatted string arrays to NumPy masked arrays.
+    The strings are assumed to be wrapped in curly braces (e.g. "{1,2,3}").
+    Any occurrence of the literal "NULL" in a cell will be masked.
+    
+    Parameters
+    ----------
+    column_data : numpy.ndarray
+        A 1D NumPy array of strings. Each element is a formatted array like "{1,2,3}".
+        
+    Returns
+    -------
+    numpy.ma.MaskedArray
+        A masked array where "NULL" entries are masked.
+    """
+    # Remove curly braces (but do not remove "NULL")
+    clean_data = np.char.replace(column_data.astype(str), "{", "")
+    clean_data = np.char.replace(clean_data, "}", "")
+    
+    # Split each string by comma into a list of items (with possible surrounding whitespace)
+    split_arrays = np.char.split(clean_data, ",")
+    
+    # --- Determine type by scanning for a first non-"NULL" value ---
+    first_value = None
+    for row in split_arrays:
+        for item in row:
+            item_str = item.strip()
+            if item_str != "NULL":
+                first_value = item_str
+                break
+        if first_value is not None:
+            break
+            
+    # If no non-NULL value is found, default to a masked object array.
+    if first_value is None:
+        data = [np.array(row) for row in split_arrays]
+        mask = [np.full(len(row), True, dtype=bool) for row in split_arrays]
+        return np.ma.masked_array(data, mask=mask)
+    
+    # Try to determine numeric type.
+    # (If first_value consists solely of digits, we'll assume integer.
+    #  Otherwise, if it can be converted to float, we'll use float.
+    #  Else, we default to string.)
+    is_integer = first_value.isdigit()
+    is_float = False
+    if not is_integer:
+        try:
+            float(first_value)
+            is_float = True
+        except Exception:
+            pass
+
+    # Prepare lists to store converted rows and corresponding masks.
+    data_list = []
+    mask_list = []
+    
+    # Conversion helper functions
+    def convert_item(item, conv):
+        item = item.strip()
+        if item == "NULL":
+            return None, True
+        else:
+            return conv(item), False
+    
+    if is_integer:
+        conv_func = int
+        dtype = np.int64
+    elif is_float:
+        conv_func = float
+        dtype = np.float64
+    else:
+        conv_func = lambda x: x
+        dtype = object
+
+    # Process each row
+    for row in split_arrays:
+        row_vals = []
+        row_mask = []
+        for item in row:
+            val, is_mask = convert_item(item, conv_func)
+            # For masked numeric values, we insert a dummy (0 or 0.0) value.
+            if is_mask:
+                if dtype in (np.int64, np.float64):
+                    row_vals.append(0)
+                else:
+                    row_vals.append("")
+            else:
+                row_vals.append(val)
+            row_mask.append(is_mask)
+        # Convert row to an array of the target dtype.
+        row_arr = np.array(row_vals, dtype=dtype)
+        data_list.append(row_arr)
+        mask_list.append(np.array(row_mask, dtype=bool))
+    
+    # Create and return a masked array.
+    return np.ma.masked_array(data_list, mask=mask_list)
+
+def format_result_table(tab):
+    if tab is None or len(tab) == 0:
+        return None
+    
+    for col in tab.colnames:
+        if len(tab[col]) == 0:
+            continue
+        if not "<U" in str(tab[col].dtype):
+            continue
+        
+        if "{" in tab[col][0]:
+            tab[col] = vectorized_string_to_masked_array(tab[col])
+
+    return tab