infrahub-server 1.6.2__py3-none-any.whl → 1.7.0__py3-none-any.whl

infrahub/actions/tasks.py

@@ -211,8 +211,10 @@ async def configure_action_rules(
     service: InfrahubServices,
 ) -> None:
     await setup_triggers_specific(
-        gatherer=gather_trigger_action_rules, trigger_type=TriggerType.ACTION_TRIGGER_RULE, db=service.database
-    )  # type: ignore[misc]
+        gatherer=gather_trigger_action_rules,  # type: ignore[arg-type]
+        trigger_type=TriggerType.ACTION_TRIGGER_RULE,
+        db=service.database,
+    )
 
 
 async def _get_targets(

infrahub/api/exceptions.py

@@ -4,12 +4,12 @@ from infrahub.exceptions import Error
 class QueryValidationError(Error):
     HTTP_CODE = 400
 
-    def __init__(self, message: str):
+    def __init__(self, message: str) -> None:
         self.message = message
 
 
 class SchemaNotValidError(Error):
     HTTP_CODE = 422
 
-    def __init__(self, message: str):
+    def __init__(self, message: str) -> None:
         self.message = message

infrahub/api/schema.py

@@ -366,6 +366,7 @@ async def load_schema(
                 diff=result.diff,
                 limit=result.diff.all,
                 update_db=True,
+                user_id=account_session.account_id,
             )
             branch.update_schema_hash()
             log.info("Schema has been updated", branch=branch.name, hash=branch.active_schema_hash.main)
@@ -375,7 +376,7 @@ async def load_schema(
                 branch.is_isolated = True
                 log.info("Branch converted to isolated mode because the schema has changed", branch=branch.name)
 
-            await branch.save(db=dbt)
+            await branch.save(db=dbt, user_id=account_session.account_id)
             updated_branch = registry.schema.get_schema_branch(name=branch.name)
             updated_hash = updated_branch.get_hash()
 
@@ -387,6 +388,7 @@ async def load_schema(
             new_schema=candidate_schema,
             previous_schema=origin_schema,
             migrations=result.migrations,
+            user_id=account_session.account_id,
         )
         migration_error_msgs = await service.workflow.execute_workflow(
             workflow=SCHEMA_APPLY_MIGRATION,

infrahub/artifacts/tasks.py

@@ -14,6 +14,7 @@ async def create(model: CheckArtifactCreate) -> ValidatorConclusion:
     await add_tags(branches=[model.branch_name], nodes=[model.target_id])
 
     client = get_client()
+    client.request_context = model.context.to_request_context()
 
     validator = await client.get(kind=InfrahubKind.ARTIFACTVALIDATOR, id=model.validator_id, include=["checks"])
 

infrahub/auth.py

@@ -73,8 +73,8 @@ async def authenticate_with_password(
 ) -> models.UserToken:
     selected_branch = await registry.get_branch(db=db, branch=branch)
 
-    response: list[CoreGenericAccount] = await NodeManager.query(
-        schema=InfrahubKind.GENERICACCOUNT,
+    response = await NodeManager.query(
+        schema=CoreGenericAccount,
         db=db,
         branch=selected_branch,
         filters={"name__value": credentials.username},

infrahub/cli/db.py

@@ -65,6 +65,7 @@ def get_timestamp_string() -> str:
 if TYPE_CHECKING:
     from infrahub.cli.context import CliContext
     from infrahub.core.migrations.shared import MigrationTypes
+    from infrahub.core.root import Root
     from infrahub.database import InfrahubDatabase
     from infrahub.database.index import IndexManagerBase
 
@@ -93,12 +94,40 @@ def callback() -> None:
     """
 
 
+async def do_migrate(
+    db: InfrahubDatabase,
+    root_node: Root,
+    check: bool = False,
+    migration_number: int | None = None,
+) -> None:
+    """Core migration logic that can be called independently of CLI.
+
+    Args:
+        db: The database connection.
+        root_node: The root node containing the current graph version.
+        check: If True, only check which migrations need to run without applying them.
+        migration_number: If provided, run only this specific migration.
+    """
+    migrations = await detect_migration_to_run(
+        current_graph_version=root_node.graph_version, migration_number=migration_number
+    )
+
+    if check or not migrations:
+        return
+
+    await migrate_database(
+        db=db, migrations=migrations, initialize=True, update_graph_version=(migration_number is None)
+    )
+
+
 @app.command(name="migrate")
 async def migrate_cmd(
     ctx: typer.Context,
     check: bool = typer.Option(False, help="Check the state of the database without applying the migrations."),
     config_file: str = typer.Argument("infrahub.toml", envvar="INFRAHUB_CONFIG"),
-    migration_number: int | None = typer.Option(None, help="Apply a specific migration by number"),
+    migration_number: int | None = typer.Option(
+        None, help="Apply a specific migration by number, regardless of current database version"
+    ),
 ) -> None:
     """Check the current format of the internal graph and apply the necessary migrations"""
     logging.getLogger("infrahub").setLevel(logging.WARNING)
@@ -111,14 +140,7 @@ async def migrate_cmd(
     dbdriver = await context.init_db(retry=1)
 
     root_node = await get_root_node(db=dbdriver)
-    migrations = await detect_migration_to_run(
-        current_graph_version=root_node.graph_version, migration_number=migration_number
-    )
-
-    if check or not migrations:
-        return
-
-    await migrate_database(db=dbdriver, migrations=migrations, initialize=True)
+    await do_migrate(db=dbdriver, root_node=root_node, check=check, migration_number=migration_number)
 
     await dbdriver.close()
 
@@ -292,18 +314,17 @@ async def detect_migration_to_run(
         migration = get_migration_by_number(migration_number)
         migrations.append(migration)
         if current_graph_version > migration.minimum_version:
+            get_migration_console().log(f"Migration {migration_number} will be re-applied.")
+        else:
             get_migration_console().log(
-                f"Migration {migration_number} already applied. To apply again, run the command without the --check flag."
+                f"Migration {migration_number} needs to be applied. Run `infrahub db migrate` to apply all outstanding migrations."
             )
-            return []
-        get_migration_console().log(
-            f"Migration {migration_number} needs to be applied. Run `infrahub db migrate` to apply all outstanding migrations."
-        )
-    else:
-        migrations.extend(await get_graph_migrations(current_graph_version=current_graph_version))
-        if not migrations:
-            get_migration_console().log(f"Database up-to-date (v{current_graph_version}), no migration to execute.")
-            return []
+            return migrations
+
+    migrations.extend(await get_graph_migrations(current_graph_version=current_graph_version))
+    if not migrations:
+        get_migration_console().log(f"Database up-to-date (v{current_graph_version}), no migration to execute.")
+        return []
 
     get_migration_console().log(
         f"Database needs to be updated (v{current_graph_version} -> v{GRAPH_VERSION}), {len(migrations)} migrations pending"
@@ -312,7 +333,10 @@ async def detect_migration_to_run(
 
 
 async def migrate_database(
-    db: InfrahubDatabase, migrations: Sequence[MigrationTypes], initialize: bool = False
+    db: InfrahubDatabase,
+    migrations: Sequence[MigrationTypes],
+    initialize: bool = False,
+    update_graph_version: bool = True,
 ) -> bool:
     """Apply the latest migrations to the database, this function will print the status directly in the console.
 
@@ -322,6 +346,7 @@ async def migrate_database(
         db: The database object.
         migrations: Sequence of migrations to apply.
         initialize: Whether to initialize the registry before running migrations.
+        update_graph_version: Whether to update the graph version after each migration.
     """
     if not migrations:
         return True
@@ -339,8 +364,9 @@ async def migrate_database(
             validation_result = await migration.validate_migration(db=db)
             if validation_result.success:
                 get_migration_console().log(f"Migration: {migration.name} {SUCCESS_BADGE}")
-                root_node.graph_version = migration.minimum_version + 1
-                await root_node.save(db=db)
+                if update_graph_version:
+                    root_node.graph_version = migration.minimum_version + 1
+                    await root_node.save(db=db)
 
         if not execution_result.success or (validation_result and not validation_result.success):
             get_migration_console().log(f"Migration: {migration.name} {FAILED_BADGE}")
@@ -603,17 +629,17 @@ RETURN vertices, edges
     edge_path = export_dir / Path("edges.csv")
     edge_path.touch(exist_ok=True)
 
-    graph_node_schemas = [GraphNodeProperties, GraphRelationshipProperties, GraphAttributeProperties]
+    graph_node_schemas = (GraphNodeProperties, GraphRelationshipProperties, GraphAttributeProperties)
     graph_vertex_properties = set()
-    for graph_schema in graph_node_schemas:
-        for field_name, field_info in graph_schema.model_fields.items():
+    for graph_node_schema in graph_node_schemas:
+        for field_name, field_info in graph_node_schema.model_fields.items():
             property_name = field_info.alias or field_name
             graph_vertex_properties.add(property_name)
 
-    graph_edge_schemas = [GraphRelationshipIsPartOf, GraphRelationshipDefault]
+    graph_edge_schemas = (GraphRelationshipIsPartOf, GraphRelationshipDefault)
     graph_edge_properties = set()
-    for graph_schema in graph_edge_schemas:
-        for field_name, field_info in graph_schema.model_fields.items():
+    for graph_edge_schema in graph_edge_schemas:
+        for field_name, field_info in graph_edge_schema.model_fields.items():
             property_name = field_info.alias or field_name
             graph_edge_properties.add(property_name)
 

infrahub/computed_attribute/gather.py

@@ -7,9 +7,9 @@ from prefect import task
 from prefect.cache_policies import NONE
 from prefect.logging import get_run_logger
 
-from infrahub.core.constants import InfrahubKind
 from infrahub.core.manager import NodeManager
 from infrahub.core.protocols import CoreGenericRepository, CoreGraphQLQuery
+from infrahub.core.protocols import CoreTransformPython as CoreTransformPythonNode
 from infrahub.core.registry import registry
 from infrahub.database import InfrahubDatabase  # noqa: TC001  needed for prefect flow
 from infrahub.git.utils import get_repositories_commit_per_branch
@@ -24,7 +24,6 @@ from .models import (
 )
 
 if TYPE_CHECKING:
-    from infrahub.core.protocols import CoreTransformPython as CoreTransformPythonNode
     from infrahub.git.models import RepositoryData
 
 
@@ -48,9 +47,9 @@ async def gather_python_transform_attributes(
     if not transform_names:
         return []
 
-    transforms: list[CoreTransformPythonNode] = await NodeManager.query(
+    transforms = await NodeManager.query(
         db=db,
-        schema=InfrahubKind.TRANSFORMPYTHON,
+        schema=CoreTransformPythonNode,
         branch=branch_name,
         fields={"id": None, "name": None, "repository": None, "query": None},
         filters={"name__values": transform_names},

infrahub/computed_attribute/tasks.py

@@ -42,8 +42,10 @@ mutation UpdateAttribute(
     $kind: String!,
     $attribute: String!,
     $value: String!
+    $context_account_id: String!
   ) {
   InfrahubUpdateComputedAttribute(
+    context: {account: {id: $context_account_id}},
     data: {id: $id, attribute: $attribute, value: $value, kind: $kind}
   ) {
     ok
@@ -62,7 +64,7 @@ async def process_transform(
     object_id: str,
     computed_attribute_name: str,  # noqa: ARG001
     computed_attribute_kind: str,  # noqa: ARG001
-    context: InfrahubContext,  # noqa: ARG001
+    context: InfrahubContext,
     updated_fields: list[str] | None = None,  # noqa: ARG001
 ) -> None:
     await add_tags(branches=[branch_name], nodes=[object_id])
@@ -103,7 +105,7 @@ async def process_transform(
             name=transform.repository.peer.name.value,
             repository_kind=str(transform.repository.peer.typename),
             commit=repo_node.commit.value,
-        )  # type: ignore[misc]
+        )
 
         data = await client.query_gql_query(
             name=transform.query.id,
@@ -120,11 +122,17 @@ async def process_transform(
             location=f"{transform.file_path.value}::{transform.class_name.value}",
             data=data,
             convert_query_response=transform.convert_query_response.value,
-        )  # type: ignore[misc]
+        )  # type: ignore[call-overload]
 
         await client.execute_graphql(
             query=UPDATE_ATTRIBUTE,
-            variables={"id": object_id, "kind": node_kind, "attribute": attribute_name, "value": transformed_data},
+            variables={
+                "id": object_id,
+                "kind": node_kind,
+                "attribute": attribute_name,
+                "value": transformed_data,
+                "context_account_id": context.account.account_id,
+            },
             branch_name=branch_name,
         )
 
@@ -168,6 +176,7 @@ async def computed_attribute_jinja2_update_value(
     node_kind: str,
     attribute_name: str,
     template: Jinja2Template,
+    context: InfrahubContext,
 ) -> None:
     log = get_run_logger()
     client = get_client()
@@ -182,7 +191,13 @@ async def computed_attribute_jinja2_update_value(
     try:
         await client.execute_graphql(
             query=UPDATE_ATTRIBUTE,
-            variables={"id": obj.node_id, "kind": node_kind, "attribute": attribute_name, "value": value},
+            variables={
+                "id": obj.node_id,
+                "kind": node_kind,
+                "attribute": attribute_name,
+                "value": value,
+                "context_account_id": context.account.account_id,
+            },
             branch_name=branch_name,
         )
         log.info(f"Updating computed attribute {node_kind}.{attribute_name}='{value}' ({obj.node_id})")
@@ -202,7 +217,7 @@ async def process_jinja2(
     object_id: str,
     computed_attribute_name: str,
     computed_attribute_kind: str,
-    context: InfrahubContext,  # noqa: ARG001
+    context: InfrahubContext,
     updated_fields: list[str] | None = None,
 ) -> None:
     log = get_run_logger()
@@ -258,6 +273,7 @@ async def process_jinja2(
                 node_kind=node_schema.kind,
                 attribute_name=computed_macro.attribute.name,
                 template=jinja_template,
+                context=context,
             )
 
         _ = [response async for _, response in batch.execute()]
@@ -462,6 +478,7 @@ async def query_transform_targets(
                         "object_id": subscriber.object_id,
                         "computed_attribute_name": computed_attribute.name,
                         "computed_attribute_kind": subscriber.kind,
+                        "context": context,
                     },
                 )
 

infrahub/config.py

@@ -207,6 +207,14 @@ class MainSettings(BaseSettings):
     def convert_to_path(cls, value: Path | str) -> Path:
         return Path(value) if isinstance(value, str) else value
 
+    @property
+    def infrahub_address(self) -> str:
+        """This is the address that the Prefect worker will use to connect to Infrahub API."""
+        if self.internal_address:
+            return self.internal_address
+
+        raise InitializationError()
+
 
 class FileSystemStorageSettings(BaseSettings):
     # Make variable lookup case-sensitive to avoid fetching $PATH value

infrahub/constants/enums.py

@@ -0,0 +1,12 @@
+from enum import StrEnum
+
+
+class OrderDirection(StrEnum):
+    ASC = "ASC"
+    DESC = "DESC"
+
+
+class OrderByField(StrEnum):
+    ID = "id"
+    CREATED_AT = "created_at"
+    UPDATED_AT = "updated_at"

infrahub/core/account.py

@@ -54,11 +54,15 @@ class AccountGlobalPermissionQuery(Query):
     name: str = "account_global_permissions"
     type: QueryType = QueryType.READ
 
-    def __init__(self, account_id: str, **kwargs: Any):
+    def __init__(
+        self,
+        account_id: str,
+        branch: Branch | None = None,
+        branch_agnostic: bool = False,
+    ) -> None:
+        super().__init__(branch=branch, branch_agnostic=branch_agnostic)
         self.account_id = account_id
-        super().__init__(**kwargs)
 
-    async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:  # noqa: ARG002
         self.params["account_id"] = self.account_id
 
         branch_filter, branch_params = self.branch.get_query_filter_path(
@@ -66,7 +70,6 @@ class AccountGlobalPermissionQuery(Query):
         )
         self.params.update(branch_params)
 
-        # ruff: noqa: E501
         query = """
         MATCH (account:%(generic_account_node)s)
         WHERE account.uuid = $account_id
@@ -164,7 +167,7 @@ class AccountObjectPermissionQuery(Query):
     name: str = "account_object_permissions"
     type: QueryType = QueryType.READ
 
-    def __init__(self, account_id: str, **kwargs: Any):
+    def __init__(self, account_id: str, **kwargs: Any) -> None:
         self.account_id = account_id
         super().__init__(**kwargs)
 
@@ -293,7 +296,7 @@ class AccountObjectPermissionQuery(Query):
 
 
 async def fetch_permissions(account_id: str, db: InfrahubDatabase, branch: Branch) -> AssignedPermissions:
-    query1 = await AccountGlobalPermissionQuery.init(db=db, branch=branch, account_id=account_id, branch_agnostic=True)
+    query1 = AccountGlobalPermissionQuery(branch=branch, account_id=account_id, branch_agnostic=True)
     await query1.execute(db=db)
     global_permissions = query1.get_permissions()
 
@@ -308,7 +311,7 @@ class AccountRoleGlobalPermissionQuery(Query):
     name: str = "account_role_global_permissions"
     type: QueryType = QueryType.READ
 
-    def __init__(self, role_id: str, **kwargs: Any):
+    def __init__(self, role_id: str, **kwargs: Any) -> None:
         self.role_id = role_id
         super().__init__(**kwargs)
 
@@ -394,7 +397,7 @@ class AccountRoleObjectPermissionQuery(Query):
     name: str = "account_role_object_permissions"
     type: QueryType = QueryType.READ
 
-    def __init__(self, role_id: str, **kwargs: Any):
+    def __init__(self, role_id: str, **kwargs: Any) -> None:
         self.role_id = role_id
         super().__init__(**kwargs)
 
@@ -515,7 +518,7 @@ class AccountTokenValidateQuery(Query):
     name: str = "account_token_validate"
     type: QueryType = QueryType.READ
 
-    def __init__(self, token: str, **kwargs: Any):
+    def __init__(self, token: str, **kwargs: Any) -> None:
         self.token = token
         super().__init__(**kwargs)