pytrilogy 0.0.3.6__tar.gz → 0.0.3.7__tar.gz

{pytrilogy-0.0.3.6/pytrilogy.egg-info → pytrilogy-0.0.3.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: pytrilogy
-Version: 0.0.3.6
+Version: 0.0.3.7
 Summary: Declarative, typed query language that compiles to SQL.
 Home-page: 
 Author: 

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.7/pytrilogy.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: pytrilogy
-Version: 0.0.3.6
+Version: 0.0.3.7
 Summary: Declarative, typed query language that compiles to SQL.
 Home-page: 
 Author: 

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.7}/pytrilogy.egg-info/SOURCES.txt

@@ -103,6 +103,7 @@ trilogy/dialect/base.py
 trilogy/dialect/bigquery.py
 trilogy/dialect/common.py
 trilogy/dialect/config.py
+trilogy/dialect/dataframe.py
 trilogy/dialect/duckdb.py
 trilogy/dialect/enums.py
 trilogy/dialect/postgres.py

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.7}/trilogy/__init__.py

@@ -4,6 +4,6 @@ from trilogy.dialect.enums import Dialects
 from trilogy.executor import Executor
 from trilogy.parser import parse
 
-__version__ = "0.0.3.6"
+__version__ = "0.0.3.7"
 
 __all__ = ["parse", "Executor", "Dialects", "Environment", "CONFIG"]

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.7}/trilogy/core/functions.py

@@ -503,32 +503,52 @@ FUNCTION_REGISTRY: dict[FunctionType, FunctionConfig] = {
         arg_count=1,
     ),
     FunctionType.ADD: FunctionConfig(
-        valid_inputs={DataType.INTEGER, DataType.FLOAT, DataType.NUMBER},
+        valid_inputs={
+            DataType.INTEGER,
+            DataType.FLOAT,
+            DataType.NUMBER,
+            DataType.NUMERIC,
+        },
         output_purpose=Purpose.PROPERTY,
         output_type=DataType.INTEGER,
         arg_count=InfiniteFunctionArgs,
     ),
     FunctionType.SUBTRACT: FunctionConfig(
-        valid_inputs={DataType.INTEGER, DataType.FLOAT, DataType.NUMBER},
+        valid_inputs={
+            DataType.INTEGER,
+            DataType.FLOAT,
+            DataType.NUMBER,
+            DataType.NUMERIC,
+        },
         output_purpose=Purpose.PROPERTY,
         output_type=DataType.INTEGER,
         arg_count=InfiniteFunctionArgs,
     ),
     FunctionType.MULTIPLY: FunctionConfig(
-        valid_inputs={DataType.INTEGER, DataType.FLOAT, DataType.NUMBER},
+        valid_inputs={
+            DataType.INTEGER,
+            DataType.FLOAT,
+            DataType.NUMBER,
+            DataType.NUMERIC,
+        },
         output_purpose=Purpose.PROPERTY,
         output_type=DataType.INTEGER,
         arg_count=InfiniteFunctionArgs,
     ),
     FunctionType.DIVIDE: FunctionConfig(
-        valid_inputs={DataType.INTEGER, DataType.FLOAT, DataType.NUMBER},
+        valid_inputs={
+            DataType.INTEGER,
+            DataType.FLOAT,
+            DataType.NUMBER,
+            DataType.NUMERIC,
+        },
         output_purpose=Purpose.PROPERTY,
         output_type=DataType.INTEGER,
         arg_count=InfiniteFunctionArgs,
     ),
     FunctionType.MOD: FunctionConfig(
         valid_inputs=[
-            {DataType.INTEGER, DataType.FLOAT, DataType.NUMBER},
+            {DataType.INTEGER, DataType.FLOAT, DataType.NUMBER, DataType.NUMERIC},
             {DataType.INTEGER},
         ],
         output_purpose=Purpose.PROPERTY,
@@ -537,7 +557,7 @@ FUNCTION_REGISTRY: dict[FunctionType, FunctionConfig] = {
     ),
     FunctionType.ROUND: FunctionConfig(
         valid_inputs=[
-            {DataType.INTEGER, DataType.FLOAT, DataType.NUMBER},
+            {DataType.INTEGER, DataType.FLOAT, DataType.NUMBER, DataType.NUMERIC},
             {DataType.INTEGER},
         ],
         output_purpose=Purpose.PROPERTY,

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.7}/trilogy/core/models/datasource.py

@@ -117,6 +117,12 @@ class Datasource(HasUUID, Namespaced, BaseModel):
     where: Optional[WhereClause] = None
     non_partial_for: Optional[WhereClause] = None
 
+    @property
+    def safe_address(self) -> str:
+        if isinstance(self.address, Address):
+            return self.address.location
+        return self.address
+
     def __eq__(self, other):
         if not isinstance(other, Datasource):
             return False

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.7}/trilogy/dialect/config.py

@@ -1,3 +1,6 @@
+from pandas import DataFrame
+
+
 class DialectConfig:
     def __init__(self):
         pass
@@ -104,3 +107,9 @@ class TrinoConfig(PrestoConfig):
         if self.schema:
             return f"trino://{self.username}:{self.password}@{self.host}:{self.port}/{self.catalog}/{self.schema}"
         return f"trino://{self.username}:{self.password}@{self.host}:{self.port}/{self.catalog}"
+
+
+class DataFrameConfig(DuckDBConfig):
+    def __init__(self, dataframes: dict[str, DataFrame]):
+        super().__init__()
+        self.dataframes = dataframes

pytrilogy-0.0.3.7/trilogy/dialect/dataframe.py

@@ -0,0 +1,42 @@
+from typing import Any
+
+from pandas import DataFrame
+from sqlalchemy import text
+
+from trilogy.core.models.environment import Environment
+from trilogy.dialect.duckdb import DuckDBDialect
+from trilogy.engine import ExecutionEngine
+
+
+class DataframeDialect(DuckDBDialect):
+    pass
+
+
+class DataframeConnectionWrapper(ExecutionEngine):
+    def __init__(self, engine: ExecutionEngine, dataframes: dict[str, DataFrame]):
+        self.engine = engine
+        self.dataframes = dataframes
+        self.connection = None
+
+    def setup(self, env: Environment, connection):
+        self._register_dataframes(env, connection)
+
+    def _register_dataframes(self, env: Environment, connection):
+        for ds in env.datasources.values():
+            if ds.safe_address in self.dataframes:
+                connection.execute(
+                    text("register(:name, :df)"),
+                    {"name": ds.safe_address, "df": self.dataframes[ds.safe_address]},
+                )
+            else:
+                raise ValueError(
+                    f"Dataframe {ds.safe_address} not found in dataframes on connection config, have {self.dataframes.keys()}"
+                )
+        pass
+
+    def add_dataframe(self, name: str, df: DataFrame, connection, env: Environment):
+        self.dataframes[name] = df
+        self._register_dataframes(env, connection)
+
+    def connect(self) -> Any:
+        return self.engine.connect()

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.7}/trilogy/dialect/enums.py

@@ -16,7 +16,7 @@ def default_factory(conf: DialectConfig, config_type):
 
     if not isinstance(conf, config_type):
         raise TypeError(
-            f"Invalid dialect configuration for type {type(config_type).__name__}"
+            f"Invalid dialect configuration for type {type(config_type).__name__}, is {type(conf)}"
         )
     if conf.connect_args:
         return create_engine(
@@ -33,6 +33,7 @@ class Dialects(Enum):
     TRINO = "trino"
     POSTGRES = "postgres"
     SNOWFLAKE = "snowflake"
+    DATAFRAME = "dataframe"
 
     @classmethod
     def _missing_(cls, value):
@@ -88,6 +89,16 @@ class Dialects(Enum):
             from trilogy.dialect.config import TrinoConfig
 
             return _engine_factory(conf, TrinoConfig)
+        elif self == Dialects.DATAFRAME:
+            from trilogy.dialect.config import DataFrameConfig
+            from trilogy.dialect.dataframe import DataframeConnectionWrapper
+
+            if not conf:
+                conf = DataFrameConfig(dataframes={})
+
+            base = _engine_factory(conf, DataFrameConfig)
+
+            return DataframeConnectionWrapper(base, dataframes=conf.dataframes)
         else:
             raise ValueError(
                 f"Unsupported dialect {self} for default engine creation; create one explicitly."

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.7}/trilogy/engine.py

@@ -1,7 +1,9 @@
-from typing import Protocol
+from typing import Any, Protocol
 
 from sqlalchemy.engine import Connection, CursorResult, Engine
 
+from trilogy.core.models.environment import Environment
+
 
 class EngineResult(Protocol):
     pass
@@ -13,7 +15,7 @@ class EngineResult(Protocol):
 class EngineConnection(Protocol):
     pass
 
-    def execute(self, statement: str) -> EngineResult:
+    def execute(self, statement: str, parameters: Any | None = None) -> EngineResult:
         pass
 
 
@@ -23,6 +25,9 @@ class ExecutionEngine(Protocol):
     def connect(self) -> EngineConnection:
         pass
 
+    def setup(self, env: Environment, connection):
+        pass
+
 
 ### Begin default SQLAlchemy implementation
 class SqlAlchemyResult(EngineResult):
@@ -37,8 +42,10 @@ class SqlAlchemyConnection(EngineConnection):
     def __init__(self, connection: Connection):
         self.connection = connection
 
-    def execute(self, statement: str) -> SqlAlchemyResult:
-        return SqlAlchemyResult(self.connection.execute(statement))
+    def execute(
+        self, statement: str, parameters: Any | None = None
+    ) -> SqlAlchemyResult:
+        return SqlAlchemyResult(self.connection.execute(statement, parameters))
 
 
 class SqlAlchemyEngine(ExecutionEngine):

{pytrilogy-0.0.3.6 → pytrilogy-0.0.3.7}/trilogy/executor.py

@@ -4,7 +4,7 @@ from pathlib import Path
 from typing import Any, Generator, List, Optional, Protocol
 
 from sqlalchemy import text
-from sqlalchemy.engine import CursorResult, Engine
+from sqlalchemy.engine import CursorResult
 
 from trilogy.constants import logger
 from trilogy.core.enums import FunctionType, Granularity, IOType
@@ -33,6 +33,7 @@ from trilogy.core.statements.execute import (
 )
 from trilogy.dialect.base import BaseDialect
 from trilogy.dialect.enums import Dialects
+from trilogy.engine import ExecutionEngine
 from trilogy.hooks.base_hook import BaseHook
 from trilogy.parser import parse_text
 
@@ -71,7 +72,7 @@ class Executor(object):
     def __init__(
         self,
         dialect: Dialects,
-        engine: Engine,
+        engine: ExecutionEngine,
         environment: Optional[Environment] = None,
         hooks: List[BaseHook] | None = None,
     ):
@@ -109,9 +110,16 @@ class Executor(object):
             from trilogy.dialect.snowflake import SnowflakeDialect
 
             self.generator = SnowflakeDialect()
+        elif self.dialect == Dialects.DATAFRAME:
+            from trilogy.dialect.dataframe import DataframeDialect
+
+            self.generator = DataframeDialect()
         else:
             raise ValueError(f"Unsupported dialect {self.dialect}")
         self.connection = self.engine.connect()
+        # TODO: make generic
+        if self.dialect == Dialects.DATAFRAME:
+            self.engine.setup(self.environment, self.connection)
 
     def execute_statement(self, statement) -> Optional[CursorResult]:
         if not isinstance(