shaped 2.0.2__py3-none-any.whl → 2.0.4__py3-none-any.whl

shaped/config_builders.py

@@ -4,47 +4,50 @@ Fluent builders for constructing Shaped V2 configurations.
 This module provides a fluent API for building Engine, Table, and View configurations
 using Python objects instead of raw dictionaries or complex model instantiation.
 """
-from typing import Any, Dict, List, Optional, Union
 
-from shaped.autogen.models.engine_config_v2 import EngineConfigV2
-from shaped.autogen.models.data_config import DataConfig
-from shaped.autogen.models.data_config_interaction_table import DataConfigInteractionTable
-from shaped.autogen.models.reference_table_config import ReferenceTableConfig
-from shaped.autogen.models.query_table_config import QueryTableConfig
+from typing import Any, Dict, List, Optional, Union
 
-# Table Config Imports
-from shaped.autogen.models.postgres_table_config import PostgresTableConfig
-from shaped.autogen.models.big_query_table_config import BigQueryTableConfig
-from shaped.autogen.models.snowflake_table_config import SnowflakeTableConfig
-from shaped.autogen.models.file_table_config import FileTableConfig
+from shaped.autogen.models.ai_enrichment_view_config import AIEnrichmentViewConfig
 from shaped.autogen.models.amplitude_table_config import AmplitudeTableConfig
 from shaped.autogen.models.aws_pinpoint_table_config import AWSPinpointTableConfig
+from shaped.autogen.models.big_query_table_config import BigQueryTableConfig
 from shaped.autogen.models.clickhouse_table_config import ClickhouseTableConfig
 from shaped.autogen.models.custom_table_config import CustomTableConfig
+from shaped.autogen.models.data_config import DataConfig
+from shaped.autogen.models.data_config_interaction_table import (
+    DataConfigInteractionTable,
+)
+from shaped.autogen.models.deployment_config import DeploymentConfig
 from shaped.autogen.models.dynamo_db_table_config import DynamoDBTableConfig
+from shaped.autogen.models.engine_config_v2 import EngineConfigV2
+from shaped.autogen.models.file_table_config import FileTableConfig
 from shaped.autogen.models.iceberg_table_config import IcebergTableConfig
+
+# Other Imports
+from shaped.autogen.models.index_config import IndexConfig
 from shaped.autogen.models.kafka_table_config import KafkaTableConfig
 from shaped.autogen.models.kinesis_table_config import KinesisTableConfig
 from shaped.autogen.models.mongo_db_table_config import MongoDBTableConfig
 from shaped.autogen.models.mssql_table_config import MSSQLTableConfig
 from shaped.autogen.models.my_sql_table_config import MySQLTableConfig
+
+# Table Config Imports
+from shaped.autogen.models.postgres_table_config import PostgresTableConfig
 from shaped.autogen.models.posthog_table_config import PosthogTableConfig
+from shaped.autogen.models.query_definition import QueryDefinition
+from shaped.autogen.models.query_table_config import QueryTableConfig
+from shaped.autogen.models.rank_query_config import RankQueryConfig
 from shaped.autogen.models.redshift_table_config import RedshiftTableConfig
+from shaped.autogen.models.reference_table_config import ReferenceTableConfig
 from shaped.autogen.models.rudderstack_table_config import RudderstackTableConfig
 from shaped.autogen.models.segment_table_config import SegmentTableConfig
 from shaped.autogen.models.shopify_table_config import ShopifyTableConfig
+from shaped.autogen.models.snowflake_table_config import SnowflakeTableConfig
+from shaped.autogen.models.sql_transform_type import SQLTransformType
 
 # View Config Imports
 from shaped.autogen.models.sql_view_config import SQLViewConfig
-from shaped.autogen.models.ai_enrichment_view_config import AIEnrichmentViewConfig
-from shaped.autogen.models.sql_transform_type import SQLTransformType
-
-# Other Imports
-from shaped.autogen.models.index_config import IndexConfig
 from shaped.autogen.models.training_config import TrainingConfig
-from shaped.autogen.models.deployment_config import DeploymentConfig
-from shaped.autogen.models.query_definition import QueryDefinition
-from shaped.autogen.models.rank_query_config import RankQueryConfig
 
 # Type aliases for better type hints
 TableConfig = Union[
@@ -67,7 +70,7 @@ TableConfig = Union[
     RedshiftTableConfig,
     RudderstackTableConfig,
     SegmentTableConfig,
-    ShopifyTableConfig
+    ShopifyTableConfig,
 ]
 ViewConfig = Union[SQLViewConfig, AIEnrichmentViewConfig]
 
@@ -75,7 +78,7 @@ ViewConfig = Union[SQLViewConfig, AIEnrichmentViewConfig]
 class Engine:
     """
     Builder for Engine configurations with support for all V2 features.
-    
+
     Example:
         engine = (Engine("my-engine")
             .interactions("events")
@@ -86,11 +89,11 @@ class Engine:
             .with_deployment_config(...)
             .with_query("recommendations", "SELECT * FROM items ORDER BY score DESC"))
     """
-    
+
     def __init__(self, name: str):
         if not name:
             raise ValueError("Engine name is required")
-            
+
         self._name = name
         self._description: Optional[str] = None
         self._interactions: Optional[Union[str, Dict[str, Any]]] = None
@@ -106,187 +109,185 @@ class Engine:
         self._training_config: Dict[str, Any] = {}
         self._deployment_config: Dict[str, Any] = {}
         self._queries: Dict[str, QueryDefinition] = {}
-    
+
     def description(self, description: str) -> "Engine":
         """Set the engine description."""
         self._description = description
         return self
-    
+
     def interactions(self, table: Union[str, Dict[str, Any]]) -> "Engine":
         """Set the interactions table configuration."""
         self._interactions = table
         return self
-    
+
     def users(self, table: Union[str, Dict[str, Any]]) -> "Engine":
         """Set the users table configuration."""
         self._users = table
         return self
-    
+
     def items(self, table: Union[str, Dict[str, Any]]) -> "Engine":
         """Set the items table configuration."""
         self._items = table
         return self
-    
+
     def schedule(self, schedule: str) -> "Engine":
         """Set the data refresh schedule."""
         self._schedule = schedule
         return self
-    
+
     def with_schema_override(self, schema: Dict[str, Any]) -> "Engine":
         """Override the default schema detection with custom schema."""
         self._schema_override = schema
         return self
-    
+
     def with_compute(self, compute: Dict[str, Any]) -> "Engine":
         """Configure compute resources for data processing."""
         self._compute = compute
         return self
-    
+
     def add_filter(self, filter_config: Dict[str, Any]) -> "Engine":
         """
         Add a filter configuration.
-        
+
         Args:
             filter_config: Dictionary containing filter configuration
-            
+
         Returns:
             The Engine instance for method chaining
         """
         self._filters.append(filter_config)
         return self
-        
+
     def add_filters(self, *filter_configs: Dict[str, Any]) -> "Engine":
         """
         Add multiple filter configurations at once.
-        
+
         Args:
             *filter_configs: One or more filter configurations to add
-            
+
         Returns:
             The Engine instance for method chaining
         """
         self._filters.extend(filter_configs)
         return self
-    
+
     def add_reference_table(self, name: str, config: Dict[str, Any]) -> "Engine":
         """
         Add a reference table configuration.
-        
+
         Args:
             name: Name of the reference table
             config: Configuration dictionary for the reference table
-            
+
         Returns:
             The Engine instance for method chaining
         """
         self._reference_tables[name] = config
         return self
-        
+
     def add_reference_tables(self, **tables: Dict[str, Any]) -> "Engine":
         """
         Add multiple reference tables at once using keyword arguments.
-        
+
         Args:
             **tables: Keyword arguments where keys are table names and values are configs
-            
+
         Returns:
             The Engine instance for method chaining
         """
         self._reference_tables.update(tables)
         return self
-    
+
     def with_tags(self, **tags: str) -> "Engine":
         """Add tags to the engine configuration."""
         if self._tags is None:
             self._tags = {}
         self._tags.update(tags)
         return self
-    
+
     def with_index_config(self, config: IndexConfig) -> "Engine":
         """Configure index settings."""
         self._index_config = config
         return self
-    
+
     def with_training_config(
         self,
         data_split_strategy: Optional[str] = None,
         data_split_params: Optional[Dict[str, Any]] = None,
-        **kwargs
+        **kwargs,
     ) -> "Engine":
         """Configure model training settings."""
         if data_split_strategy:
             self._training_config["data_split"] = {
                 "strategy": data_split_strategy,
-                **(data_split_params or {})
+                **(data_split_params or {}),
             }
         self._training_config.update(kwargs)
         return self
-    
+
     def with_deployment_config(self, **kwargs) -> "Engine":
         """Configure deployment settings."""
         self._deployment_config.update(kwargs)
         return self
-    
+
     def with_query(self, name: str, query: Union[str, Dict[str, Any]]) -> "Engine":
         """Add a named query definition."""
         if isinstance(query, str):
             # Convert SQL string to a RankQueryConfig
-            query_def = QueryDefinition(
-                query=RankQueryConfig(type="rank", sql=query)
-            )
+            query_def = QueryDefinition(query=RankQueryConfig(type="rank", sql=query))
         else:
             query_def = QueryDefinition(**query)
-        
+
         self._queries[name] = query_def
         return self
-    
+
     def validate(self) -> List[str]:
         """
         Validate the current configuration.
-        
+
         Returns:
             List of error messages, empty if configuration is valid
         """
         errors = []
-        
+
         if not self._interactions:
             errors.append("Interactions table is required")
-            
+
         # Validate reference tables
         for name, config in self._reference_tables.items():
             if not isinstance(config, dict) or not config:
                 errors.append(f"Invalid configuration for reference table '{name}'")
-                
+
         # Validate queries
         for name, query in self._queries.items():
-            if not query or not getattr(query, 'query', None):
+            if not query or not getattr(query, "query", None):
                 errors.append(f"Invalid query definition for query '{name}'")
-                
+
         return errors
-        
+
     def is_valid(self) -> bool:
         """
         Check if the current configuration is valid.
-        
+
         Returns:
             True if configuration is valid, False otherwise
         """
         return not bool(self.validate())
-        
+
     def build(self) -> EngineConfigV2:
         """
         Build the EngineConfigV2 object with all configured settings.
-        
+
         Returns:
             EngineConfigV2: The built engine configuration
-            
+
         Raises:
             ValueError: If the configuration is invalid
         """
         # Validate before building
         if errors := self.validate():
             raise ValueError(f"Invalid configuration: {', '.join(errors)}")
-        
+
         # Build data configuration
         data_config = DataConfig(
             interaction_table=self._create_data_table(self._interactions),
@@ -296,49 +297,51 @@ class Engine:
             schema=self._schema_override,
             compute=self._compute,
             filters=self._filters if self._filters else None,
-            reference_tables=self._reference_tables if self._reference_tables else None
+            reference_tables=self._reference_tables if self._reference_tables else None,
         )
-        
+
         # Build engine config
         config = EngineConfigV2(
             name=self._name,
             description=self._description,
             data=data_config,
             version="v2",
-            tags=self._tags
+            tags=self._tags,
         )
-        
+
         # Add optional configurations
         if self._index_config:
             config.index = self._index_config
-        
+
         if self._training_config:
             config.training = TrainingConfig(**self._training_config)
-        
+
         if self._deployment_config:
             config.deployment = DeploymentConfig(**self._deployment_config)
-        
+
         if self._queries:
             config.queries = self._queries
-        
+
         return config
 
-    def _create_data_table(self, table: Union[str, Dict[str, Any], Any]) -> Optional[DataConfigInteractionTable]:
+    def _create_data_table(
+        self, table: Union[str, Dict[str, Any], Any]
+    ) -> Optional[DataConfigInteractionTable]:
         """
         Create a DataConfigInteractionTable from input.
-        
+
         Args:
             table: Can be a table name (str), a table config dict, or an existing config object.
-            
+
         Returns:
             A properly configured DataConfigInteractionTable instance, or None if input is None.
-            
+
         Raises:
             ValueError: If the input cannot be converted to a valid table configuration.
         """
         if table is None:
             return None
-            
+
         try:
             if isinstance(table, str):
                 # It's a reference to a table by name
@@ -346,7 +349,7 @@ class Engine:
                 # Create the DataConfigInteractionTable with the config
                 result = DataConfigInteractionTable(actual_instance=config)
                 return result
-            
+
             if isinstance(table, dict):
                 # It's a config dictionary
                 if table.get("type") == "query":
@@ -357,13 +360,15 @@ class Engine:
                     config = ReferenceTableConfig(**table)
                     result = DataConfigInteractionTable(actual_instance=config)
                     return result
-                    
+
             # Check if it's already a DataConfigInteractionTable
-            if hasattr(table, 'actual_instance') and hasattr(table, 'to_dict'):
+            if hasattr(table, "actual_instance") and hasattr(table, "to_dict"):
                 return table
-                
-            raise ValueError(f"Unsupported table configuration type: {type(table).__name__}")
-            
+
+            raise ValueError(
+                f"Unsupported table configuration type: {type(table).__name__}"
+            )
+
         except Exception as e:
             raise ValueError(f"Failed to create data table: {str(e)}")
 
@@ -372,7 +377,7 @@ class Table:
     """
     Factory methods for Table configurations.
     """
-    
+
     @staticmethod
     def Postgres(
         name: str,
@@ -384,7 +389,7 @@ class Table:
         table: str,
         replication_key: str,
         description: Optional[str] = None,
-        schema: Optional[str] = None, # database_schema
+        schema: Optional[str] = None,  # database_schema
         ssl_args: Optional[Dict[str, Any]] = None,
     ) -> PostgresTableConfig:
         """Create a Postgres table configuration."""
@@ -453,7 +458,7 @@ class Table:
             replication_key=replication_key,
             description=description,
         )
-    
+
     @staticmethod
     def File(
         name: str,
@@ -484,7 +489,7 @@ class Table:
     ) -> MySQLTableConfig:
         """
         Create a MySQL table configuration.
-        
+
         Args:
             name: Name of the table configuration
             host: MySQL server hostname or IP
@@ -496,7 +501,7 @@ class Table:
             replication_key: Column name to use for change data capture
             description: Optional description of the table
             ssl_args: Optional SSL configuration parameters
-            
+
         Returns:
             MySQLTableConfig: Configured MySQL table configuration
         """
@@ -531,7 +536,7 @@ class Table:
     ) -> MSSQLTableConfig:
         """
         Create a Microsoft SQL Server table configuration.
-        
+
         Args:
             name: Name of the table configuration
             host: SQL Server hostname or IP
@@ -544,7 +549,7 @@ class Table:
             replication_key: Column name to use for change data capture
             description: Optional description of the table
             ssl_args: Optional SSL configuration parameters
-            
+
         Returns:
             MSSQLTableConfig: Configured SQL Server table configuration
         """
@@ -576,7 +581,7 @@ class Table:
     ) -> MongoDBTableConfig:
         """
         Create a MongoDB table configuration.
-        
+
         Args:
             name: Name of the table configuration
             connection_uri: MongoDB connection URI
@@ -585,7 +590,7 @@ class Table:
             replication_key: Field name to use for change data capture
             description: Optional description of the table
             **kwargs: Additional MongoDB connection parameters
-            
+
         Returns:
             MongoDBTableConfig: Configured MongoDB table configuration
         """
@@ -596,7 +601,7 @@ class Table:
             collection=collection,
             replication_key=replication_key,
             description=description,
-            **kwargs
+            **kwargs,
         )
 
     @staticmethod
@@ -615,7 +620,7 @@ class Table:
     ) -> RedshiftTableConfig:
         """
         Create an Amazon Redshift table configuration.
-        
+
         Args:
             name: Name of the table configuration
             host: Redshift cluster endpoint
@@ -628,7 +633,7 @@ class Table:
             replication_key: Column name to use for change data capture
             description: Optional description of the table
             ssl_args: Optional SSL configuration parameters
-            
+
         Returns:
             RedshiftTableConfig: Configured Redshift table configuration
         """
@@ -648,11 +653,12 @@ class Table:
             kwargs.update(ssl_args)
         return RedshiftTableConfig(**kwargs)
 
+
 class View:
     """
     Factory methods for View configurations.
     """
-    
+
     @staticmethod
     def SQL(
         name: str,
@@ -664,11 +670,15 @@ class View:
         Create a SQL view configuration.
         """
         transform_type = "MATERIALIZED_VIEW" if materialized else "VIEW"
-        
+
         return SQLViewConfig(
             name=name,
             sql_query=query,
-            sql_transform_type=SQLTransformType(transform_type) if isinstance(transform_type, str) else transform_type,
+            sql_transform_type=(
+                SQLTransformType(transform_type)
+                if isinstance(transform_type, str)
+                else transform_type
+            ),
             description=description,
         )