Spaces:
Running
Running
| from typing import Any, Dict, Iterable, List, Optional, Tuple, Union | |
| from cassandra.cluster import ResponseFuture | |
| from cassio.table.cql import DELETE_CQL_TEMPLATE, SELECT_CQL_TEMPLATE, CQLOpType | |
| from cassio.table.table_types import ColumnSpecType, RowType, normalize_type_desc | |
| from cassio.table.utils import ( | |
| call_wrapped_async, | |
| handle_multicolumn_packing, | |
| handle_multicolumn_unpacking, | |
| ) | |
| from .base_table import BaseTableMixin | |
| PARTITION_ID_TYPE = Union[Any, Tuple[Any]] | |
| class ClusteredMixin(BaseTableMixin): | |
| def __init__( | |
| self, | |
| *pargs: Any, | |
| partition_id_type: Union[str, List[str]] = ["TEXT"], | |
| partition_id: Optional[PARTITION_ID_TYPE] = None, | |
| ordering_in_partition: Union[str, List[str]] = "ASC", | |
| **kwargs: Any, | |
| ) -> None: | |
| self.partition_id_type = normalize_type_desc(partition_id_type) | |
| self.partition_id = partition_id | |
| if isinstance(ordering_in_partition, str): | |
| self.ordering_in_partition = ordering_in_partition.upper() | |
| else: | |
| self.ordering_in_partition = [ | |
| ordering.upper() for ordering in ordering_in_partition | |
| ] | |
| super().__init__(*pargs, **kwargs) | |
| def _schema_pk(self) -> List[ColumnSpecType]: | |
| if len(self.partition_id_type) == 1: | |
| return [ | |
| ("partition_id", self.partition_id_type[0]), | |
| ] | |
| else: | |
| return [ | |
| (f"partition_id_{pk_i}", pk_typ) | |
| for pk_i, pk_typ in enumerate(self.partition_id_type) | |
| ] | |
| def _schema_cc(self) -> List[ColumnSpecType]: | |
| return self._schema_row_id() | |
| def _delete_partition( | |
| self, is_async: bool, partition_id: Optional[PARTITION_ID_TYPE] = None | |
| ) -> Union[None, ResponseFuture]: | |
| _partition_id = self.partition_id if partition_id is None else partition_id | |
| # | |
| _pid_dict = handle_multicolumn_unpacking( | |
| {"partition_id": _partition_id}, | |
| "partition_id", | |
| [col for col, _ in self._schema_pk()], | |
| ) | |
| ( | |
| rest_kwargs, | |
| where_clause_blocks, | |
| delete_cql_vals, | |
| ) = self._extract_where_clause_blocks(_pid_dict) | |
| assert rest_kwargs == {} | |
| where_clause = "WHERE " + " AND ".join(where_clause_blocks) | |
| delete_cql = DELETE_CQL_TEMPLATE.format( | |
| where_clause=where_clause, | |
| ) | |
| if is_async: | |
| return self.execute_cql_async( | |
| delete_cql, args=delete_cql_vals, op_type=CQLOpType.WRITE | |
| ) | |
| else: | |
| self.execute_cql(delete_cql, args=delete_cql_vals, op_type=CQLOpType.WRITE) | |
| return None | |
| def delete_partition( | |
| self, partition_id: Optional[PARTITION_ID_TYPE] = None | |
| ) -> None: | |
| self._delete_partition(is_async=False, partition_id=partition_id) | |
| return None | |
| def delete_partition_async( | |
| self, partition_id: Optional[PARTITION_ID_TYPE] = None | |
| ) -> ResponseFuture: | |
| return self._delete_partition(is_async=True, partition_id=partition_id) | |
| async def adelete_partition( | |
| self, partition_id: Optional[PARTITION_ID_TYPE] = None | |
| ) -> None: | |
| await call_wrapped_async(self.delete_partition_async, partition_id=partition_id) | |
| def _normalize_kwargs(self, args_dict: Dict[str, Any]) -> Dict[str, Any]: | |
| # if partition id provided in call, takes precedence over instance value | |
| arg_pid = args_dict.get("partition_id") | |
| instance_pid = self.partition_id | |
| _partition_id = instance_pid if arg_pid is None else arg_pid | |
| new_args_dict0 = { | |
| **{"partition_id": _partition_id}, | |
| **args_dict, | |
| } | |
| # in case of multicolumn-key schema, do the tuple unpacking: | |
| new_args_dict = handle_multicolumn_unpacking( | |
| new_args_dict0, | |
| "partition_id", | |
| [col for col, _ in self._schema_pk()], | |
| ) | |
| return super()._normalize_kwargs(new_args_dict) | |
| def _normalize_row(self, raw_row: Any) -> Dict[str, Any]: | |
| pre_normalized = super()._normalize_row(raw_row) | |
| repacked_row = handle_multicolumn_packing( | |
| unpacked_row=pre_normalized, | |
| key_name="partition_id", | |
| unpacked_keys=[col for col, _ in self._schema_pk()], | |
| ) | |
| return repacked_row | |
| def _get_get_partition_cql( | |
| self, | |
| partition_id: Optional[PARTITION_ID_TYPE] = None, | |
| n: Optional[int] = None, | |
| **kwargs: Any, | |
| ) -> Tuple[str, Tuple[Any, ...]]: | |
| _partition_id = self.partition_id if partition_id is None else partition_id | |
| # | |
| # TODO: work on a columns: Optional[List[str]] = None | |
| # (but with nuanced handling of the column-magic we have here) | |
| columns = None | |
| if columns is None: | |
| columns_desc = "*" | |
| else: | |
| # TODO: handle translations here? | |
| # columns_desc = ", ".join(columns) | |
| raise NotImplementedError("Column selection is not implemented.") | |
| # WHERE can admit other sources (e.g. medata if the corresponding mixin) | |
| # so we escalate to standard WHERE-creation route and reinject the partition | |
| n_kwargs = self._normalize_kwargs( | |
| { | |
| **{"partition_id": _partition_id}, | |
| **kwargs, | |
| } | |
| ) | |
| ( | |
| rest_kwargs, | |
| where_clause_blocks, | |
| select_cql_vals, | |
| ) = self._extract_where_clause_blocks(n_kwargs) | |
| # check for exhaustion: | |
| assert rest_kwargs == {} | |
| where_clause = "WHERE " + " AND ".join(where_clause_blocks) | |
| where_cql_vals = list(select_cql_vals) | |
| # | |
| if n is None: | |
| limit_clause = "" | |
| limit_cql_vals = [] | |
| else: | |
| limit_clause = "LIMIT %s" | |
| limit_cql_vals = [n] | |
| # | |
| select_cql = SELECT_CQL_TEMPLATE.format( | |
| columns_desc=columns_desc, | |
| where_clause=where_clause, | |
| limit_clause=limit_clause, | |
| ) | |
| get_p_cql_vals = tuple(where_cql_vals + limit_cql_vals) | |
| return select_cql, get_p_cql_vals | |
| def get_partition( | |
| self, | |
| partition_id: Optional[PARTITION_ID_TYPE] = None, | |
| n: Optional[int] = None, | |
| **kwargs: Any, | |
| ) -> Iterable[RowType]: | |
| select_cql, get_p_cql_vals = self._get_get_partition_cql( | |
| partition_id, n, **kwargs | |
| ) | |
| return ( | |
| self._normalize_row(raw_row) | |
| for raw_row in self.execute_cql( | |
| select_cql, | |
| args=get_p_cql_vals, | |
| op_type=CQLOpType.READ, | |
| ) | |
| ) | |
| def get_partition_async( | |
| self, | |
| partition_id: Optional[PARTITION_ID_TYPE] = None, | |
| n: Optional[int] = None, | |
| **kwargs: Any, | |
| ) -> ResponseFuture: | |
| raise NotImplementedError("Asynchronous reads are not supported.") | |
| async def aget_partition( | |
| self, | |
| partition_id: Optional[PARTITION_ID_TYPE] = None, | |
| n: Optional[int] = None, | |
| **kwargs: Any, | |
| ) -> Iterable[RowType]: | |
| select_cql, get_p_cql_vals = self._get_get_partition_cql( | |
| partition_id, n, **kwargs | |
| ) | |
| return ( | |
| self._normalize_row(raw_row) | |
| for raw_row in await self.aexecute_cql( | |
| select_cql, | |
| args=get_p_cql_vals, | |
| op_type=CQLOpType.READ, | |
| ) | |
| ) | |