pytrilogy 0.0.3.94__py3-none-any.whl → 0.0.3.95__py3-none-any.whl

trilogy/core/env_processor.py

@@ -20,7 +20,8 @@ def add_concept(
     if node_name in seen:
         return
     seen.add(node_name)
-    g.add_node(concept)
+    g.concepts[node_name] = concept
+    g.add_node(node_name)
     if concept.concept_arguments:
         for source in concept.concept_arguments:
             if not isinstance(source, BuildConcept):
@@ -28,9 +29,10 @@ def add_concept(
                     f"Invalid non-build concept {source} passed into graph generation from {concept}"
                 )
             generic = get_default_grain_concept(source, default_concept_graph)
+            generic_node = concept_to_node(generic)
             add_concept(generic, g, concept_mapping, default_concept_graph, seen)
 
-            g.add_edge(generic, node_name)
+            g.add_edge(generic_node, node_name, fast=True)
     for ps_address in concept.pseudonyms:
         if ps_address not in concept_mapping:
             raise SyntaxError(f"Concept {concept} has invalid pseudonym {ps_address}")
@@ -44,8 +46,8 @@ def add_concept(
             continue
         if pseudonym_node.split("@")[0] == node_name.split("@")[0]:
             continue
-        g.add_edge(pseudonym_node, node_name)
-        g.add_edge(node_name, pseudonym_node)
+        g.add_edge(pseudonym_node, node_name, fast=True)
+        g.add_edge(node_name, pseudonym_node, fast=True)
         g.pseudonyms.add((pseudonym_node, node_name))
         g.pseudonyms.add((node_name, pseudonym_node))
         add_concept(pseudonym, g, concept_mapping, default_concept_graph, seen)
@@ -82,20 +84,27 @@ def generate_adhoc_graph(
 
     for dataset in datasources:
         node = datasource_to_node(dataset)
-        g.add_node(dataset, type="datasource", datasource=dataset)
+        g.add_datasource_node(node, dataset)
         for concept in dataset.concepts:
+            cnode = concept_to_node(concept)
+            g.concepts[cnode] = concept
+            g.add_node(cnode)
             if restrict_to_listed:
-                if concept_to_node(concept) not in g.nodes:
+                if cnode not in g.nodes:
                     continue
-            g.add_edge(node, concept)
-            g.add_edge(concept, node)
+            g.add_edge(node, cnode, fast=True)
+            g.add_edge(cnode, node, fast=True)
             # if there is a key on a table at a different grain
             # add an FK edge to the canonical source, if it exists
             # for example, order ID on order product table
             default = get_default_grain_concept(concept, default_concept_graph)
+
             if concept != default:
-                g.add_edge(concept, default)
-                g.add_edge(default, concept)
+                dcnode = concept_to_node(default)
+                g.concepts[dcnode] = default
+                g.add_node(dcnode)
+                g.add_edge(cnode, dcnode, fast=True)
+                g.add_edge(dcnode, cnode, fast=True)
     return g
 
 

trilogy/core/environment_helpers.py

@@ -169,6 +169,112 @@ def generate_key_concepts(concept: Concept, environment: Environment):
         environment.add_concept(new_concept, add_derived=False)
 
 
+def remove_date_concepts(concept: Concept, environment: Environment):
+    """Remove auto-generated date-related concepts for the given concept"""
+    date_suffixes = ["month", "year", "quarter", "day", "day_of_week"]
+    grain_suffixes = ["month_start", "year_start"]
+
+    for suffix in date_suffixes + grain_suffixes:
+        address = concept.address + f".{suffix}"
+        if address in environment.concepts:
+            derived_concept = environment.concepts[address]
+            # Only remove if it was auto-derived from this concept
+            if (
+                derived_concept.metadata
+                and derived_concept.metadata.concept_source
+                == ConceptSource.AUTO_DERIVED
+                and derived_concept.keys
+                and concept.address in derived_concept.keys
+            ):
+                environment.remove_concept(address)
+
+
+def remove_datetime_concepts(concept: Concept, environment: Environment):
+    """Remove auto-generated datetime-related concepts for the given concept"""
+    datetime_suffixes = ["date", "hour", "minute", "second"]
+
+    for suffix in datetime_suffixes:
+        address = concept.address + f".{suffix}"
+        if address in environment.concepts:
+            derived_concept = environment.concepts[address]
+            # Only remove if it was auto-derived from this concept
+            if (
+                derived_concept.metadata
+                and derived_concept.metadata.concept_source
+                == ConceptSource.AUTO_DERIVED
+                and derived_concept.keys
+                and concept.address in derived_concept.keys
+            ):
+                environment.remove_concept(address)
+
+
+def remove_key_concepts(concept: Concept, environment: Environment):
+    """Remove auto-generated key-related concepts for the given concept"""
+    key_suffixes = ["count"]
+
+    for suffix in key_suffixes:
+        address = concept.address + f".{suffix}"
+        if address in environment.concepts:
+            derived_concept = environment.concepts[address]
+            if (
+                derived_concept.metadata
+                and derived_concept.metadata.concept_source
+                == ConceptSource.AUTO_DERIVED
+            ):
+                environment.remove_concept(address)
+
+
+def remove_struct_concepts(concept: Concept, environment: Environment):
+    """Remove auto-generated struct field concepts for the given concept"""
+    if not isinstance(concept.datatype, StructType):
+        return
+
+    target_namespace = (
+        environment.namespace + "." + concept.name
+        if environment.namespace and environment.namespace != DEFAULT_NAMESPACE
+        else concept.name
+    )
+
+    # Get all concepts in the target namespace that were auto-derived
+    concepts_to_remove = []
+    for address, derived_concept in environment.concepts.items():
+        if (
+            derived_concept.namespace == target_namespace
+            and derived_concept.metadata
+            and derived_concept.metadata.concept_source == ConceptSource.AUTO_DERIVED
+            and isinstance(derived_concept.lineage, Function)
+            and derived_concept.lineage.operator == FunctionType.ATTR_ACCESS
+            and len(derived_concept.lineage.arguments) >= 1
+            and derived_concept.lineage.arguments[0] == concept.reference
+        ):
+            concepts_to_remove.append(address)
+
+    for address in concepts_to_remove:
+        environment.remove_concept(address)
+
+
+def remove_related_concepts(concept: Concept, environment: Environment):
+    """Remove all auto-generated concepts that were derived from the given concept"""
+
+    # Remove key-related concepts
+    if concept.purpose == Purpose.KEY:
+        remove_key_concepts(concept, environment)
+
+    # Remove datatype-specific concepts
+    if concept.datatype == DataType.DATE:
+        remove_date_concepts(concept, environment)
+    elif concept.datatype == DataType.DATETIME:
+        remove_date_concepts(concept, environment)
+        remove_datetime_concepts(concept, environment)
+    elif concept.datatype == DataType.TIMESTAMP:
+        remove_date_concepts(concept, environment)
+        remove_datetime_concepts(concept, environment)
+
+    # Remove struct field concepts
+    if isinstance(concept.datatype, StructType):
+        remove_struct_concepts(concept, environment)
+
+
 def generate_related_concepts(
     concept: Concept,
     environment: Environment,
@@ -183,6 +289,7 @@ def generate_related_concepts(
     if concept.datatype == DataType.DATE and add_derived:
         generate_date_concepts(concept, environment)
     elif concept.datatype == DataType.DATETIME and add_derived:
+
         generate_date_concepts(concept, environment)
         generate_datetime_concepts(concept, environment)
     elif concept.datatype == DataType.TIMESTAMP and add_derived:
@@ -203,6 +310,10 @@ def generate_related_concepts(
                 ),
                 lineage=AttrAccess([concept.reference, key], environment=environment),
                 grain=concept.grain,
+                metadata=Metadata(
+                    concept_source=ConceptSource.AUTO_DERIVED,
+                ),
+                keys=concept.keys,
             )
             environment.add_concept(auto, meta=meta)
             if isinstance(value, Concept):

trilogy/core/exceptions.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Sequence
 
 
 class UndefinedConceptException(Exception):
@@ -24,6 +24,26 @@ class NoDatasourceException(UnresolvableQueryException):
     pass
 
 
+class ModelValidationError(Exception):
+    def __init__(
+        self,
+        message,
+        children: Sequence["ModelValidationError"] | None = None,
+        **kwargs
+    ):
+        super().__init__(self, message, **kwargs)
+        self.message = message
+        self.children = children
+
+
+class DatasourceModelValidationError(ModelValidationError):
+    pass
+
+
+class ConceptModelValidationError(ModelValidationError):
+    pass
+
+
 class AmbiguousRelationshipResolutionException(UnresolvableQueryException):
     def __init__(self, message, parents: List[set[str]]):
         super().__init__(self, message)

trilogy/core/functions.py

@@ -380,7 +380,12 @@ FUNCTION_REGISTRY: dict[FunctionType, FunctionConfig] = {
     ),
     FunctionType.CURRENT_DATETIME: FunctionConfig(
         output_purpose=Purpose.CONSTANT,
-        output_type=DataType.DATE,
+        output_type=DataType.DATETIME,
+        arg_count=0,
+    ),
+    FunctionType.CURRENT_TIMESTAMP: FunctionConfig(
+        output_purpose=Purpose.CONSTANT,
+        output_type=DataType.TIMESTAMP,
         arg_count=0,
     ),
     FunctionType.BOOL: FunctionConfig(

trilogy/core/graph_models.py

@@ -93,40 +93,14 @@ class ReferenceGraph(nx.DiGraph):
     def add_node(self, node_for_adding, fast: bool = False, **attr):
         if fast:
             return super().add_node(node_for_adding, **attr)
-        if isinstance(node_for_adding, BuildConcept):
-            node_name = concept_to_node(node_for_adding)
-            self.concepts[node_name] = node_for_adding
-        elif isinstance(node_for_adding, BuildDatasource):
-            node_name = datasource_to_node(node_for_adding)
-            self.datasources[node_name] = node_for_adding
-        else:
-            node_name = node_for_adding
-            if attr.get("datasource"):
-                self.datasources[node_name] = attr["datasource"]
+        node_name = node_for_adding
+        if attr.get("datasource"):
+            self.datasources[node_name] = attr["datasource"]
         super().add_node(node_name, **attr)
 
+    def add_datasource_node(self, node_name, datasource):
+        self.datasources[node_name] = datasource
+        super().add_node(node_name, datasource=datasource)
+
     def add_edge(self, u_of_edge, v_of_edge, fast: bool = False, **attr):
-        if fast:
-            return super().add_edge(u_of_edge, v_of_edge, **attr)
-        if isinstance(u_of_edge, BuildConcept):
-            orig = u_of_edge
-            u_of_edge = concept_to_node(u_of_edge)
-            if u_of_edge not in self.nodes:
-                self.add_node(orig)
-        elif isinstance(u_of_edge, BuildDatasource):
-            origd = u_of_edge
-            u_of_edge = datasource_to_node(u_of_edge)
-            if u_of_edge not in self.nodes:
-                self.add_node(origd)
-
-        if isinstance(v_of_edge, BuildConcept):
-            orig = v_of_edge
-            v_of_edge = concept_to_node(v_of_edge)
-            if v_of_edge not in self.nodes:
-                self.add_node(orig)
-        elif isinstance(v_of_edge, BuildDatasource):
-            origd = v_of_edge
-            v_of_edge = datasource_to_node(v_of_edge)
-            if v_of_edge not in self.nodes:
-                self.add_node(origd)
-        super().add_edge(u_of_edge, v_of_edge)
+        return super().add_edge(u_of_edge, v_of_edge, **attr)

trilogy/core/internal.py

@@ -64,4 +64,22 @@ DEFAULT_CONCEPTS = {
         granularity=Granularity.SINGLE_ROW,
         derivation=Derivation.CONSTANT,
     ),
+    "label": Concept(
+        name="label",
+        namespace=INTERNAL_NAMESPACE,
+        datatype=DataType.STRING,
+        purpose=Purpose.KEY,
+        grain=Grain(),
+        granularity=Granularity.SINGLE_ROW,
+        derivation=Derivation.CONSTANT,
+    ),
+    "expected": Concept(
+        name="expected_value",
+        namespace=INTERNAL_NAMESPACE,
+        datatype=DataType.STRING,
+        purpose=Purpose.KEY,
+        grain=Grain(),
+        granularity=Granularity.SINGLE_ROW,
+        derivation=Derivation.CONSTANT,
+    ),
 }

trilogy/core/models/core.py

@@ -103,6 +103,9 @@ class DataType(Enum):
     def data_type(self):
         return self
 
+    def __str__(self) -> str:
+        return self.name
+
 
 class TraitDataType(BaseModel):
     type: DataType | NumericType | StructType | ArrayType | MapType

trilogy/core/models/environment.py

@@ -571,6 +571,34 @@ class Environment(BaseModel):
 
         return concept
 
+    def remove_concept(
+        self,
+        concept: Concept | str,
+    ) -> bool:
+        if self.frozen:
+            raise FrozenEnvironmentException(
+                "Environment is frozen, cannot remove concepts"
+            )
+        if isinstance(concept, Concept):
+            address = concept.address
+            c_instance = concept
+        else:
+            address = concept
+            c_instance_check = self.concepts.get(address)
+            if not c_instance_check:
+                return False
+            c_instance = c_instance_check
+        from trilogy.core.environment_helpers import remove_related_concepts
+
+        remove_related_concepts(c_instance, self)
+        if address in self.concepts:
+            del self.concepts[address]
+            return True
+        if address in self.alias_origin_lookup:
+            del self.alias_origin_lookup[address]
+
+        return False
+
     def add_datasource(
         self,
         datasource: Datasource,

trilogy/core/models/execute.py

@@ -23,6 +23,7 @@ from trilogy.core.constants import CONSTANT_DATASET
 from trilogy.core.enums import (
     ComparisonOperator,
     Derivation,
+    FunctionClass,
     FunctionType,
     JoinType,
     Modifier,
@@ -375,6 +376,12 @@ class CTE(BaseModel):
                 return check_is_not_in_group(c.lineage.content)
             if c.derivation == Derivation.CONSTANT:
                 return True
+            if (
+                c.purpose == Purpose.CONSTANT
+                and isinstance(c.lineage, BuildFunction)
+                and c.lineage.operator in FunctionClass.AGGREGATE_FUNCTIONS.value
+            ):
+                return True
             if c.purpose == Purpose.METRIC:
                 return True
 

trilogy/core/processing/node_generators/select_merge_node.py

@@ -123,7 +123,7 @@ def create_pruned_concept_graph(
         common: set[BuildConcept] = set.intersection(
             *[set(x.output_concepts) for x in ds_list]
         )
-        g.add_node(node_address, datasource=ds_list)
+        g.add_datasource_node(node_address, ds_list)
         for c in common:
             cnode = concept_to_node(c)
             g.add_edge(node_address, cnode)
@@ -214,7 +214,7 @@ def create_pruned_concept_graph(
     relevant = set(relevant_concepts + relevent_datasets)
     for edge in orig_g.edges():
         if edge[0] in relevant and edge[1] in relevant:
-            g.add_edge(edge[0], edge[1])
+            g.add_edge(edge[0], edge[1], fast=True)
     # if we have no ds nodes at all, for non constant, we can't find it
     if not any([n.startswith("ds~") for n in g.nodes]):
         logger.info(

trilogy/core/query_processor.py

@@ -40,8 +40,8 @@ from trilogy.core.statements.author import (
     PersistStatement,
     SelectStatement,
 )
-from trilogy.core.statements.common import MaterializedDataset
 from trilogy.core.statements.execute import (
+    MaterializedDataset,
     ProcessedCopyStatement,
     ProcessedQuery,
     ProcessedQueryPersist,
@@ -567,4 +567,5 @@ def process_query(
         base=root_cte,
         hidden_columns=set([x for x in statement.hidden_components]),
         local_concepts=statement.local_concepts,
+        locally_derived=statement.locally_derived,
     )

trilogy/core/statements/author.py

@@ -12,6 +12,7 @@ from trilogy.core.enums import (
     IOType,
     Modifier,
     ShowCategory,
+    ValidationScope,
 )
 from trilogy.core.models.author import (
     AggregateWrapper,
@@ -147,11 +148,13 @@ class SelectStatement(HasUUID, SelectTypeMixin, BaseModel):
                     continue
                 if CONFIG.parsing.select_as_definition and not environment.frozen:
                     if x.concept.address not in environment.concepts:
-                        environment.add_concept(x.content.output)
+                        environment.add_concept(x.content.output, add_derived=False)
                     elif x.concept.address in environment.concepts:
                         version = environment.concepts[x.concept.address]
                         if version.metadata.concept_source == ConceptSource.SELECT:
-                            environment.add_concept(x.content.output, force=True)
+                            environment.add_concept(
+                                x.content.output, force=True, add_derived=False
+                            )
                 x.content.output = x.content.output.set_select_grain(
                     output.grain, environment
                 )
@@ -378,6 +381,13 @@ class MultiSelectStatement(HasUUID, SelectTypeMixin, BaseModel):
             output = output.union(select.hidden_components)
         return output
 
+    @property
+    def locally_derived(self) -> set[str]:
+        locally_derived: set[str] = set([x.address for x in self.derived_concepts])
+        for select in self.selects:
+            locally_derived = locally_derived.union(select.locally_derived)
+        return locally_derived
+
 
 class RowsetDerivationStatement(HasUUID, BaseModel):
     name: str
@@ -428,8 +438,13 @@ class PersistStatement(HasUUID, BaseModel):
         return self.datasource.address
 
 
+class ValidateStatement(BaseModel):
+    scope: ValidationScope
+    targets: Optional[List[str]] = None  # list of identifiers
+
+
 class ShowStatement(BaseModel):
-    content: SelectStatement | PersistStatement | ShowCategory
+    content: SelectStatement | PersistStatement | ValidateStatement | ShowCategory
 
 
 class Limit(BaseModel):

trilogy/core/statements/common.py

@@ -4,7 +4,6 @@ from pydantic import BaseModel, Field
 
 from trilogy.core.enums import IOType
 from trilogy.core.models.author import ConceptRef, HavingClause, WhereClause
-from trilogy.core.models.datasource import Address, Datasource
 
 
 class CopyQueryMixin(BaseModel):
@@ -12,15 +11,6 @@ class CopyQueryMixin(BaseModel):
     target_type: IOType
 
 
-class MaterializedDataset(BaseModel):
-    address: Address
-
-
-class PersistQueryMixin(BaseModel):
-    output_to: MaterializedDataset
-    datasource: Datasource
-
-
 class SelectTypeMixin(BaseModel):
     where_clause: Union["WhereClause", None] = Field(default=None)
     having_clause: Union["HavingClause", None] = Field(default=None)

trilogy/core/statements/execute.py

@@ -1,49 +1,104 @@
-from typing import Annotated, List, Optional, Union
+from dataclasses import dataclass, field
+from typing import List, Optional, Union
 
-from pydantic import BaseModel, Field
-from pydantic.functional_validators import PlainValidator
-
-from trilogy.core.models.author import ConceptRef
+from trilogy.core.enums import IOType, ValidationScope
+from trilogy.core.models.author import ConceptRef, HavingClause, WhereClause
 from trilogy.core.models.build import (
     BuildConcept,
     BuildDatasource,
     BuildOrderBy,
 )
-from trilogy.core.models.environment import EnvironmentConceptDict, validate_concepts
+from trilogy.core.models.datasource import Address, Datasource
+from trilogy.core.models.environment import EnvironmentConceptDict
 from trilogy.core.models.execute import CTE, UnionCTE
-from trilogy.core.statements.common import CopyQueryMixin, PersistQueryMixin
 
 
-class ProcessedQuery(BaseModel):
+@dataclass
+class CopyQueryMixin:
+    target: str
+    target_type: IOType
+
+
+@dataclass
+class MaterializedDataset:
+    address: Address
+
+
+@dataclass
+class PersistQueryMixin:
+    output_to: MaterializedDataset
+    datasource: Datasource
+
+
+@dataclass
+class SelectTypeMixin:
+    where_clause: Union["WhereClause", None] = field(default=None)
+    having_clause: Union["HavingClause", None] = field(default=None)
+
+    @property
+    def output_components(self) -> List[ConceptRef]:
+        raise NotImplementedError
+
+
+@dataclass
+class ProcessedQuery:
     output_columns: List[ConceptRef]
     ctes: List[CTE | UnionCTE]
     base: CTE | UnionCTE
-    hidden_columns: set[str] = Field(default_factory=set)
+    hidden_columns: set[str] = field(default_factory=set)
     limit: Optional[int] = None
     order_by: Optional[BuildOrderBy] = None
-    local_concepts: Annotated[
-        EnvironmentConceptDict, PlainValidator(validate_concepts)
-    ] = Field(default_factory=EnvironmentConceptDict)
+    local_concepts: EnvironmentConceptDict = field(
+        default_factory=EnvironmentConceptDict
+    )
+    locally_derived: set[str] = field(default_factory=set)
 
 
+@dataclass
 class ProcessedQueryPersist(ProcessedQuery, PersistQueryMixin):
     pass
 
 
+@dataclass
 class ProcessedCopyStatement(ProcessedQuery, CopyQueryMixin):
     pass
 
 
-class ProcessedRawSQLStatement(BaseModel):
+@dataclass
+class ProcessedRawSQLStatement:
     text: str
 
 
-class ProcessedStaticValueOutput(BaseModel):
+@dataclass
+class ProcessedValidateStatement:
+    scope: ValidationScope
+    targets: Optional[List[str]]
+
+
+@dataclass
+class ProcessedStaticValueOutput:
     values: List[dict]
 
 
-class ProcessedShowStatement(BaseModel):
+@dataclass
+class ProcessedShowStatement:
     output_columns: List[ConceptRef]
     output_values: List[
-        Union[BuildConcept, BuildDatasource, ProcessedQuery, ProcessedStaticValueOutput]
+        Union[
+            BuildConcept,
+            BuildDatasource,
+            ProcessedQuery,
+            ProcessedValidateStatement,
+            ProcessedStaticValueOutput,
+        ]
     ]
+
+
+PROCESSED_STATEMENT_TYPES = (
+    ProcessedCopyStatement
+    | ProcessedQuery
+    | ProcessedRawSQLStatement
+    | ProcessedQueryPersist
+    | ProcessedShowStatement
+    | ProcessedValidateStatement
+)