moose-lib 0.6.81__tar.gz → 0.6.82__tar.gz

{moose_lib-0.6.81 → moose_lib-0.6.82}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: moose_lib
-Version: 0.6.81
+Version: 0.6.82
 Home-page: https://www.fiveonefour.com/moose
 Author: Fiveonefour Labs Inc.
 Author-email: support@fiveonefour.com

{moose_lib-0.6.81 → moose_lib-0.6.82}/moose_lib/__init__.py

@@ -26,3 +26,5 @@ from .blocks import (
 )
 from .data_models import Key, AggregateFunction, StringToEnumMixin
 from .commons import Logger
+
+from .query_builder import *

{moose_lib-0.6.81 → moose_lib-0.6.82}/moose_lib/data_models.py

@@ -147,6 +147,11 @@ class Column(BaseModel):
     default: str | None = None
     annotations: list[Tuple[str, Any]] = []
 
+    def to_expr(self):
+        # Lazy import to avoid circular dependency at import time
+        from .query_builder import ColumnRef
+        return ColumnRef(self)
+
 
 def py_type_to_column_type(t: type, mds: list[Any]) -> Tuple[bool, list[Any], DataType]:
     # handle Annotated[Optional[Annotated[...], ...]

{moose_lib-0.6.81 → moose_lib-0.6.82}/moose_lib/main.py

@@ -27,6 +27,7 @@ from .data_models import Column
 from .config.runtime import RuntimeClickHouseConfig
 
 from moose_lib.commons import EnhancedJSONEncoder
+from .query_builder import Query
 
 
 @dataclass
@@ -184,7 +185,29 @@ class QueryClient:
     def __call__(self, input, variables):
         return self.execute(input, variables)
 
-    def execute(self, input, variables, row_type: Type[BaseModel] = None):
+    def execute(self, input: Union[str, Query], variables = None, row_type: Type[BaseModel] = None):
+        """
+        Execute a query.
+
+        - If `input` is a `Query`, do not supply `variables`.
+        - If `input` is a `str` intended for `string.Formatter` interpolation, `variables` must be a dict
+          mapping placeholder names to values.
+
+        Args:
+            input: Either a `Query` object or a `Formatter`-style SQL template string.
+            variables: Dict used to fill a `Formatter` string; must be omitted when `input` is a `Query`.
+            row_type: Optional Pydantic model class to map result rows into.
+
+        Returns:
+            A list of row dicts, or a list of `row_type` instances if provided.
+        """
+        if isinstance(input, Query):
+            if variables is not None:
+                raise ValueError("Do not supply variables when you provide Query")
+            sql, params = input.to_sql_and_params()
+            print(f"[QueryClient] | Query: {sql}")
+            return self.execute_raw(sql, params, row_type)
+
         params = {}
         values: dict[str, Any] = {}
         preview_params = {}
@@ -200,25 +223,10 @@ class QueryClient:
                     params[variable_name] = f'{{p{i}: Identifier}}'
                     values[f'p{i}'] = value.name
                 else:
-                    if isinstance(value, bool):
-                        params[variable_name] = f'{{p{i}: Bool}}'
-                        values[f'p{i}'] = value
-                    elif isinstance(value, datetime):
-                        params[variable_name] = f'{{p{i}: DateTime}}'
-                        values[f'p{i}'] = value
-                    elif isinstance(value, int):
-                        params[variable_name] = f'{{p{i}: Int64}}'
-                        values[f'p{i}'] = value
-                    elif isinstance(value, float):
-                        params[variable_name] = f'{{p{i}: Float64}}'
-                        values[f'p{i}'] = value
-                    elif isinstance(value, str):
-                        params[variable_name] = f'{{p{i}: String}}'
-                        values[f'p{i}'] = value
-                    else:
-                        print(f"unhandled type in QueryClient {type(value)}", file=sys.stderr)
-                        params[variable_name] = f'{{p{i}: String}}'
-                        values[f'p{i}'] = str(value)
+                    from moose_lib.utilities.sql import clickhouse_param_type_for_value
+                    ch_type = clickhouse_param_type_for_value(value)
+                    params[variable_name] = f'{{p{i}: {ch_type}}}'
+                    values[f'p{i}'] = value
                 preview_params[variable_name] = self._format_value_for_preview(value)
 
         clickhouse_query = input.format_map(params)

moose_lib-0.6.82/moose_lib/query_builder.py

@@ -0,0 +1,198 @@
+from datetime import datetime
+from typing import Callable, Literal
+
+import sqlglot as sg
+import sqlglot.expressions as sge
+from pydantic import BaseModel
+
+from .data_models import Column, Key
+from .dmv2 import OlapTable, IngestPipeline, IngestPipelineConfig
+from moose_lib.utilities.sql import clickhouse_param_type_for_value
+
+
+class Params:
+    def __init__(self):
+        self._counter = 0
+        self.bindings: dict[str, object] = {}
+
+    def bind(self, value: object, name: str | None = None, ch_type: str | None = None) -> sge.Expression:
+        if name is None:
+            name = f"p{self._counter}"
+            self._counter += 1
+
+        if ch_type is None:
+            ch_type = clickhouse_param_type_for_value(value)
+
+        expr = sg.parse_one(f"{{{name}: {ch_type}}}", dialect="clickhouse")
+        self.bindings[name] = value
+        return expr
+
+
+def to_column(col: Column, table_name: str | None = None) -> sge.Column:
+    col_name = getattr(col, "name")
+    table_ident = None
+    if table_name is not None:
+        table_ident = sge.Identifier(this=table_name, quoted=True)
+    elif hasattr(col, "table_name") and getattr(col, "table_name"):
+        table_ident = sge.Identifier(this=getattr(col, "table_name"), quoted=True)
+
+    return sge.Column(
+        this=sge.Identifier(this=col_name, quoted=True),
+        table=table_ident,
+    )
+
+
+type Predicate = Callable[["Query"], sge.Expression]
+
+# Order-by type annotations
+type OrderDirection = Literal["asc", "desc"]
+type OrderExpr = sge.Expression | Column | "ColumnRef"
+type OrderItem = OrderExpr | tuple[OrderExpr, OrderDirection]
+
+
+class ColumnRef:
+    def __init__(self, column: Column):
+        self._column = column
+
+    def _binary_op(self, op_name: str, value: object) -> Predicate:
+        def resolve(query: "Query") -> sge.Expression:
+            table_name = query._from_table.name if query._from_table is not None else None
+            left = to_column(self._column, table_name)
+            right = query.params.bind(value)
+            op = getattr(left, op_name)
+            return op(right)
+
+        return resolve
+
+    def eq(self, value: object) -> Predicate:
+        return self._binary_op("eq", value)
+
+    def ne(self, value: object) -> Predicate:
+        return self._binary_op("neq", value)
+
+    def lt(self, value: object) -> Predicate:
+        return self._binary_op("__lt__", value)
+
+    def le(self, value: object) -> Predicate:
+        return self._binary_op("__le__", value)
+
+    def gt(self, value: object) -> Predicate:
+        return self._binary_op("__gt__", value)
+
+    def ge(self, value: object) -> Predicate:
+        return self._binary_op("__ge__", value)
+
+    def in_(self, values: list[object]) -> Predicate:
+        def resolve(query: "Query") -> sge.Expression:
+            table_name = query._from_table.name if query._from_table is not None else None
+            left = to_column(self._column, table_name)
+            rights = [query.params.bind(v) for v in values]
+            return left.isin(*rights)
+
+        return resolve
+
+    def is_null(self) -> Predicate:
+        def resolve(query: "Query") -> sge.Expression:
+            table_name = query._from_table.name if query._from_table is not None else None
+            left = to_column(self._column, table_name)
+            return left.is_(sge.Null())
+
+        return resolve
+
+
+def col(column: Column) -> ColumnRef:
+    return ColumnRef(column)
+
+
+class Query:
+    def __init__(self):
+        self.params = Params()
+        self.inner: sge.Select = sge.Select()
+        self._from_table: OlapTable | None = None
+
+    def from_(self, table: OlapTable) -> "Query":
+        self._from_table = table
+        self.inner = self.inner.from_(table.name)
+        return self
+
+    def select(self, *cols: Column) -> "Query":
+        sge_cols = [to_column(c, self._from_table.name if self._from_table is not None else None) for c in cols]
+        self.inner = self.inner.select(*sge_cols)
+        return self
+
+    def where(self, predicate_or_expr) -> "Query":
+        if callable(predicate_or_expr):
+            expr = predicate_or_expr(self)
+        else:
+            expr = predicate_or_expr
+        self.inner = self.inner.where(expr)
+        return self
+
+    def order_by(self, *items: OrderItem) -> "Query":
+        orders: list[sge.Expression] = []
+        table_name = self._from_table.name if self._from_table is not None else None
+
+        for item in items:
+            desc = False
+            expr: sge.Expression
+
+            if isinstance(item, tuple) and len(item) == 2:
+                col_like, direction = item
+                if isinstance(direction, str):
+                    desc = direction.lower() == "desc"
+                else:
+                    raise ValueError("order_by direction must be 'asc' or 'desc'")
+                if isinstance(col_like, Column):
+                    expr = to_column(col_like, table_name)
+                elif isinstance(col_like, ColumnRef):
+                    expr = to_column(col_like._column, table_name)
+                else:
+                    expr = col_like
+            else:
+                if isinstance(item, Column):
+                    expr = to_column(item, table_name)
+                elif isinstance(item, ColumnRef):
+                    expr = to_column(item._column, table_name)
+                else:
+                    expr = item
+
+            orders.append(sge.Ordered(this=expr, desc=desc))
+
+        self.inner = self.inner.order_by(*orders)
+        return self
+
+    def to_sql(self) -> str:
+        return self.inner.sql(dialect="clickhouse")
+
+    def to_sql_and_params(self) -> tuple[str, dict[str, object]]:
+        return self.to_sql(), dict(self.params.bindings)
+
+    def limit(self, n: int) -> "Query":
+        self.inner = self.inner.limit(n)
+        return self
+
+
+def and_(*predicates_or_exprs) -> Predicate:
+    def resolve(query: "Query") -> sge.Expression:
+        exprs = [p(query) if callable(p) else p for p in predicates_or_exprs]
+        if not exprs:
+            raise ValueError("and_ requires at least one predicate")
+        combined = exprs[0]
+        for e in exprs[1:]:
+            combined = combined.and_(e)
+        return combined
+
+    return resolve
+
+
+def or_(*predicates_or_exprs) -> Predicate:
+    def resolve(query: "Query") -> sge.Expression:
+        exprs = [p(query) if callable(p) else p for p in predicates_or_exprs]
+        if not exprs:
+            raise ValueError("or_ requires at least one predicate")
+        combined = exprs[0]
+        for e in exprs[1:]:
+            combined = combined.or_(e)
+        return combined
+
+    return resolve

moose_lib-0.6.82/moose_lib/utilities/sql.py

@@ -0,0 +1,34 @@
+import sys
+
+
+def quote_identifier(name: str) -> str:
+    """Quote a ClickHouse identifier with backticks if not already quoted.
+
+    Backticks allow special characters (e.g., hyphens) in identifiers.
+    """
+    if name.startswith("`") and name.endswith("`"):
+        return name
+    return f"`{name}`"
+
+
+from datetime import datetime
+from typing import Any
+
+
+def clickhouse_param_type_for_value(value: Any) -> str:
+    """Infer ClickHouse typed parameter annotation for a Python value.
+
+    Normalized to common scalar types used in placeholders.
+    """
+    if isinstance(value, bool):
+        return "Bool"
+    if isinstance(value, int):
+        return "Int64"
+    if isinstance(value, float):
+        return "Float64"
+    if isinstance(value, datetime):
+        return "DateTime"
+    if not isinstance(value, str):
+        print(f"unhandled type {type(value)}", file=sys.stderr)
+    return "String"
+

{moose_lib-0.6.81 → moose_lib-0.6.82}/moose_lib.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: moose_lib
-Version: 0.6.81
+Version: 0.6.82
 Home-page: https://www.fiveonefour.com/moose
 Author: Fiveonefour Labs Inc.
 Author-email: support@fiveonefour.com

{moose_lib-0.6.81 → moose_lib-0.6.82}/moose_lib.egg-info/SOURCES.txt

@@ -7,6 +7,7 @@ moose_lib/data_models.py
 moose_lib/dmv2_serializer.py
 moose_lib/internal.py
 moose_lib/main.py
+moose_lib/query_builder.py
 moose_lib/query_param.py
 moose_lib.egg-info/PKG-INFO
 moose_lib.egg-info/SOURCES.txt
@@ -39,5 +40,6 @@ moose_lib/utilities/sql.py
 tests/__init__.py
 tests/conftest.py
 tests/test_moose.py
+tests/test_query_builder.py
 tests/test_redis_client.py
 tests/test_s3queue_config.py

moose_lib-0.6.82/tests/test_query_builder.py

@@ -0,0 +1,38 @@
+from datetime import datetime
+
+from moose_lib.query_builder import Query, col
+from moose_lib.dmv2 import IngestPipeline, IngestPipelineConfig
+from pydantic import BaseModel
+from moose_lib.data_models import Key
+
+
+class Bar(BaseModel):
+    primary_key: Key[str]
+    utc_timestamp: datetime
+    has_text: bool
+    text_length: int
+
+
+def test_simple_select_and_where():
+    bar_model = IngestPipeline[Bar]("Bar", IngestPipelineConfig(
+        ingest=False,
+        stream=True,
+        table=True,
+        dead_letter_queue=True
+    ))
+    bar_cols = bar_model.get_table().cols
+
+    q1 = Query().from_(bar_model.get_table()).select(bar_cols.has_text, bar_cols.text_length)
+    assert q1.to_sql() == 'SELECT "Bar"."has_text", "Bar"."text_length" FROM Bar'
+
+    q2 = (
+        Query()
+        .from_(bar_model.get_table())
+        .select(bar_cols.has_text, bar_cols.text_length)
+        .where(col(bar_cols.has_text).eq(True))
+    )
+    sql, params = q2.to_sql_and_params()
+    assert sql == 'SELECT "Bar"."has_text", "Bar"."text_length" FROM Bar WHERE "Bar"."has_text" = {p0: Bool}'
+    assert params == {"p0": True}
+
+

moose_lib-0.6.81/moose_lib/utilities/sql.py

@@ -1,10 +0,0 @@
-def quote_identifier(name: str) -> str:
-    """Quote a ClickHouse identifier with backticks if not already quoted.
-
-    Backticks allow special characters (e.g., hyphens) in identifiers.
-    """
-    if name.startswith("`") and name.endswith("`"):
-        return name
-    return f"`{name}`"
-
-