genesis-flow 1.0.0__py3-none-any.whl

mlflow/server/auth/sqlalchemy_store.py

@@ -0,0 +1,263 @@
+from typing import Optional
+
+from sqlalchemy.exc import IntegrityError, MultipleResultsFound, NoResultFound
+from sqlalchemy.orm import sessionmaker
+from werkzeug.security import check_password_hash, generate_password_hash
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import (
+    INVALID_STATE,
+    RESOURCE_ALREADY_EXISTS,
+    RESOURCE_DOES_NOT_EXIST,
+)
+from mlflow.server.auth.db import utils as dbutils
+from mlflow.server.auth.db.models import (
+    SqlExperimentPermission,
+    SqlRegisteredModelPermission,
+    SqlUser,
+)
+from mlflow.server.auth.entities import ExperimentPermission, RegisteredModelPermission, User
+from mlflow.server.auth.permissions import _validate_permission
+from mlflow.store.db.utils import _get_managed_session_maker, create_sqlalchemy_engine_with_retry
+from mlflow.utils.uri import extract_db_type_from_uri
+from mlflow.utils.validation import _validate_password, _validate_username
+
+
+class SqlAlchemyStore:
+    def init_db(self, db_uri):
+        self.db_uri = db_uri
+        self.db_type = extract_db_type_from_uri(db_uri)
+        self.engine = create_sqlalchemy_engine_with_retry(db_uri)
+        dbutils.migrate_if_needed(self.engine, "head")
+        SessionMaker = sessionmaker(bind=self.engine)
+        self.ManagedSessionMaker = _get_managed_session_maker(SessionMaker, self.db_type)
+
+    def authenticate_user(self, username: str, password: str) -> bool:
+        with self.ManagedSessionMaker() as session:
+            try:
+                user = self._get_user(session, username)
+                return check_password_hash(user.password_hash, password)
+            except MlflowException:
+                return False
+
+    def create_user(self, username: str, password: str, is_admin: bool = False) -> User:
+        _validate_username(username)
+        _validate_password(password)
+        pwhash = generate_password_hash(password)
+        with self.ManagedSessionMaker() as session:
+            try:
+                user = SqlUser(username=username, password_hash=pwhash, is_admin=is_admin)
+                session.add(user)
+                session.flush()
+                return user.to_mlflow_entity()
+            except IntegrityError as e:
+                raise MlflowException(
+                    f"User (username={username}) already exists. Error: {e}",
+                    RESOURCE_ALREADY_EXISTS,
+                ) from e
+
+    @staticmethod
+    def _get_user(session, username: str) -> SqlUser:
+        try:
+            return session.query(SqlUser).filter(SqlUser.username == username).one()
+        except NoResultFound:
+            raise MlflowException(
+                f"User with username={username} not found",
+                RESOURCE_DOES_NOT_EXIST,
+            )
+        except MultipleResultsFound:
+            raise MlflowException(
+                f"Found multiple users with username={username}",
+                INVALID_STATE,
+            )
+
+    def has_user(self, username: str) -> bool:
+        with self.ManagedSessionMaker() as session:
+            return session.query(SqlUser).filter(SqlUser.username == username).first() is not None
+
+    def get_user(self, username: str) -> User:
+        with self.ManagedSessionMaker() as session:
+            return self._get_user(session, username).to_mlflow_entity()
+
+    def list_users(self) -> list[User]:
+        with self.ManagedSessionMaker() as session:
+            users = session.query(SqlUser).all()
+            return [u.to_mlflow_entity() for u in users]
+
+    def update_user(
+        self, username: str, password: Optional[str] = None, is_admin: Optional[bool] = None
+    ) -> User:
+        with self.ManagedSessionMaker() as session:
+            user = self._get_user(session, username)
+            if password is not None:
+                pwhash = generate_password_hash(password)
+                user.password_hash = pwhash
+            if is_admin is not None:
+                user.is_admin = is_admin
+            return user.to_mlflow_entity()
+
+    def delete_user(self, username: str):
+        with self.ManagedSessionMaker() as session:
+            user = self._get_user(session, username)
+            session.delete(user)
+
+    def create_experiment_permission(
+        self, experiment_id: str, username: str, permission: str
+    ) -> ExperimentPermission:
+        _validate_permission(permission)
+        with self.ManagedSessionMaker() as session:
+            try:
+                user = self._get_user(session, username=username)
+                perm = SqlExperimentPermission(
+                    experiment_id=experiment_id, user_id=user.id, permission=permission
+                )
+                session.add(perm)
+                session.flush()
+                return perm.to_mlflow_entity()
+            except IntegrityError as e:
+                raise MlflowException(
+                    f"Experiment permission (experiment_id={experiment_id}, username={username}) "
+                    f"already exists. Error: {e}",
+                    RESOURCE_ALREADY_EXISTS,
+                )
+
+    def _get_experiment_permission(
+        self, session, experiment_id: str, username: str
+    ) -> SqlExperimentPermission:
+        try:
+            user = self._get_user(session, username=username)
+            return (
+                session.query(SqlExperimentPermission)
+                .filter(
+                    SqlExperimentPermission.experiment_id == experiment_id,
+                    SqlExperimentPermission.user_id == user.id,
+                )
+                .one()
+            )
+        except NoResultFound:
+            raise MlflowException(
+                f"Experiment permission with experiment_id={experiment_id} and "
+                f"username={username} not found",
+                RESOURCE_DOES_NOT_EXIST,
+            )
+        except MultipleResultsFound:
+            raise MlflowException(
+                f"Found multiple experiment permissions with experiment_id={experiment_id} "
+                f"and username={username}",
+                INVALID_STATE,
+            )
+
+    def get_experiment_permission(self, experiment_id: str, username: str) -> ExperimentPermission:
+        with self.ManagedSessionMaker() as session:
+            return self._get_experiment_permission(
+                session, experiment_id, username
+            ).to_mlflow_entity()
+
+    def list_experiment_permissions(self, username: str) -> list[ExperimentPermission]:
+        with self.ManagedSessionMaker() as session:
+            user = self._get_user(session, username=username)
+            perms = (
+                session.query(SqlExperimentPermission)
+                .filter(SqlExperimentPermission.user_id == user.id)
+                .all()
+            )
+            return [p.to_mlflow_entity() for p in perms]
+
+    def update_experiment_permission(
+        self, experiment_id: str, username: str, permission: str
+    ) -> ExperimentPermission:
+        _validate_permission(permission)
+        with self.ManagedSessionMaker() as session:
+            perm = self._get_experiment_permission(session, experiment_id, username)
+            perm.permission = permission
+            return perm.to_mlflow_entity()
+
+    def delete_experiment_permission(self, experiment_id: str, username: str):
+        with self.ManagedSessionMaker() as session:
+            perm = self._get_experiment_permission(session, experiment_id, username)
+            session.delete(perm)
+
+    def create_registered_model_permission(
+        self, name: str, username: str, permission: str
+    ) -> RegisteredModelPermission:
+        _validate_permission(permission)
+        with self.ManagedSessionMaker() as session:
+            try:
+                user = self._get_user(session, username=username)
+                perm = SqlRegisteredModelPermission(
+                    name=name, user_id=user.id, permission=permission
+                )
+                session.add(perm)
+                session.flush()
+                return perm.to_mlflow_entity()
+            except IntegrityError as e:
+                raise MlflowException(
+                    f"Registered model permission (name={name}, username={username}) "
+                    f"already exists. Error: {e}",
+                    RESOURCE_ALREADY_EXISTS,
+                )
+
+    def _get_registered_model_permission(
+        self, session, name: str, username: str
+    ) -> SqlRegisteredModelPermission:
+        try:
+            user = self._get_user(session, username=username)
+            return (
+                session.query(SqlRegisteredModelPermission)
+                .filter(
+                    SqlRegisteredModelPermission.name == name,
+                    SqlRegisteredModelPermission.user_id == user.id,
+                )
+                .one()
+            )
+        except NoResultFound:
+            raise MlflowException(
+                f"Registered model permission with name={name} and username={username} not found",
+                RESOURCE_DOES_NOT_EXIST,
+            )
+        except MultipleResultsFound:
+            raise MlflowException(
+                f"Found multiple registered model permissions with name={name} "
+                f"and username={username}",
+                INVALID_STATE,
+            )
+
+    def get_registered_model_permission(
+        self, name: str, username: str
+    ) -> RegisteredModelPermission:
+        with self.ManagedSessionMaker() as session:
+            return self._get_registered_model_permission(session, name, username).to_mlflow_entity()
+
+    def list_registered_model_permissions(self, username: str) -> list[RegisteredModelPermission]:
+        with self.ManagedSessionMaker() as session:
+            user = self._get_user(session, username=username)
+            perms = (
+                session.query(SqlRegisteredModelPermission)
+                .filter(SqlRegisteredModelPermission.user_id == user.id)
+                .all()
+            )
+            return [p.to_mlflow_entity() for p in perms]
+
+    def update_registered_model_permission(
+        self, name: str, username: str, permission: str
+    ) -> RegisteredModelPermission:
+        _validate_permission(permission)
+        with self.ManagedSessionMaker() as session:
+            perm = self._get_registered_model_permission(session, name, username)
+            perm.permission = permission
+            return perm.to_mlflow_entity()
+
+    def delete_registered_model_permission(self, name: str, username: str):
+        with self.ManagedSessionMaker() as session:
+            perm = self._get_registered_model_permission(session, name, username)
+            session.delete(perm)
+
+    def rename_registered_model_permissions(self, old_name: str, new_name: str):
+        with self.ManagedSessionMaker() as session:
+            perms = (
+                session.query(SqlRegisteredModelPermission)
+                .filter(SqlRegisteredModelPermission.name == old_name)
+                .all()
+            )
+            for perm in perms:
+                perm.name = new_name

mlflow/server/graphql/autogenerated_graphql_schema.py

@@ -0,0 +1,353 @@
+# GENERATED FILE. PLEASE DON'T MODIFY.
+# Run python3 ./dev/proto_to_graphql/code_generator.py to regenerate.
+import graphene
+import mlflow
+from mlflow.server.graphql.graphql_custom_scalars import LongString
+from mlflow.server.graphql.graphql_errors import ApiError
+from mlflow.utils.proto_json_utils import parse_dict
+
+
+class MlflowDeploymentJobConnectionState(graphene.Enum):
+    DEPLOYMENT_JOB_CONNECTION_STATE_UNSPECIFIED = 1
+    NOT_SET_UP = 2
+    CONNECTED = 3
+    NOT_FOUND = 4
+    REQUIRED_PARAMETERS_CHANGED = 5
+
+
+class MlflowModelVersionDeploymentJobStateDeploymentJobRunState(graphene.Enum):
+    DEPLOYMENT_JOB_RUN_STATE_UNSPECIFIED = 1
+    NO_VALID_DEPLOYMENT_JOB_FOUND = 2
+    RUNNING = 3
+    SUCCEEDED = 4
+    FAILED = 5
+    PENDING = 6
+    APPROVAL = 7
+
+
+class MlflowModelVersionStatus(graphene.Enum):
+    PENDING_REGISTRATION = 1
+    FAILED_REGISTRATION = 2
+    READY = 3
+
+
+class MlflowRunStatus(graphene.Enum):
+    RUNNING = 1
+    SCHEDULED = 2
+    FINISHED = 3
+    FAILED = 4
+    KILLED = 5
+
+
+class MlflowViewType(graphene.Enum):
+    ACTIVE_ONLY = 1
+    DELETED_ONLY = 2
+    ALL = 3
+
+
+class MlflowModelVersionDeploymentJobState(graphene.ObjectType):
+    job_id = graphene.String()
+    run_id = graphene.String()
+    job_state = graphene.Field(MlflowDeploymentJobConnectionState)
+    run_state = graphene.Field(MlflowModelVersionDeploymentJobStateDeploymentJobRunState)
+    current_task_name = graphene.String()
+
+
+class MlflowModelMetric(graphene.ObjectType):
+    key = graphene.String()
+    value = graphene.Float()
+    timestamp = LongString()
+    step = LongString()
+    dataset_name = graphene.String()
+    dataset_digest = graphene.String()
+    model_id = graphene.String()
+    run_id = graphene.String()
+
+
+class MlflowModelParam(graphene.ObjectType):
+    name = graphene.String()
+    value = graphene.String()
+
+
+class MlflowModelVersionTag(graphene.ObjectType):
+    key = graphene.String()
+    value = graphene.String()
+
+
+class MlflowModelVersion(graphene.ObjectType):
+    name = graphene.String()
+    version = graphene.String()
+    creation_timestamp = LongString()
+    last_updated_timestamp = LongString()
+    user_id = graphene.String()
+    current_stage = graphene.String()
+    description = graphene.String()
+    source = graphene.String()
+    run_id = graphene.String()
+    status = graphene.Field(MlflowModelVersionStatus)
+    status_message = graphene.String()
+    tags = graphene.List(graphene.NonNull(MlflowModelVersionTag))
+    run_link = graphene.String()
+    aliases = graphene.List(graphene.String)
+    model_id = graphene.String()
+    model_params = graphene.List(graphene.NonNull(MlflowModelParam))
+    model_metrics = graphene.List(graphene.NonNull(MlflowModelMetric))
+    deployment_job_state = graphene.Field(MlflowModelVersionDeploymentJobState)
+
+
+class MlflowSearchModelVersionsResponse(graphene.ObjectType):
+    model_versions = graphene.List(graphene.NonNull(MlflowModelVersion))
+    next_page_token = graphene.String()
+    apiError = graphene.Field(ApiError)
+
+
+class MlflowDatasetSummary(graphene.ObjectType):
+    experiment_id = graphene.String()
+    name = graphene.String()
+    digest = graphene.String()
+    context = graphene.String()
+
+
+class MlflowSearchDatasetsResponse(graphene.ObjectType):
+    dataset_summaries = graphene.List(graphene.NonNull(MlflowDatasetSummary))
+    apiError = graphene.Field(ApiError)
+
+
+class MlflowMetricWithRunId(graphene.ObjectType):
+    key = graphene.String()
+    value = graphene.Float()
+    timestamp = LongString()
+    step = LongString()
+    run_id = graphene.String()
+
+
+class MlflowGetMetricHistoryBulkIntervalResponse(graphene.ObjectType):
+    metrics = graphene.List(graphene.NonNull(MlflowMetricWithRunId))
+    apiError = graphene.Field(ApiError)
+
+
+class MlflowFileInfo(graphene.ObjectType):
+    path = graphene.String()
+    is_dir = graphene.Boolean()
+    file_size = LongString()
+
+
+class MlflowListArtifactsResponse(graphene.ObjectType):
+    root_uri = graphene.String()
+    files = graphene.List(graphene.NonNull(MlflowFileInfo))
+    next_page_token = graphene.String()
+    apiError = graphene.Field(ApiError)
+
+
+class MlflowModelOutput(graphene.ObjectType):
+    model_id = graphene.String()
+    step = LongString()
+
+
+class MlflowRunOutputs(graphene.ObjectType):
+    model_outputs = graphene.List(graphene.NonNull(MlflowModelOutput))
+
+
+class MlflowModelInput(graphene.ObjectType):
+    model_id = graphene.String()
+
+
+class MlflowDataset(graphene.ObjectType):
+    name = graphene.String()
+    digest = graphene.String()
+    source_type = graphene.String()
+    source = graphene.String()
+    schema = graphene.String()
+    profile = graphene.String()
+
+
+class MlflowInputTag(graphene.ObjectType):
+    key = graphene.String()
+    value = graphene.String()
+
+
+class MlflowDatasetInput(graphene.ObjectType):
+    tags = graphene.List(graphene.NonNull(MlflowInputTag))
+    dataset = graphene.Field(MlflowDataset)
+
+
+class MlflowRunInputs(graphene.ObjectType):
+    dataset_inputs = graphene.List(graphene.NonNull(MlflowDatasetInput))
+    model_inputs = graphene.List(graphene.NonNull(MlflowModelInput))
+
+
+class MlflowRunTag(graphene.ObjectType):
+    key = graphene.String()
+    value = graphene.String()
+
+
+class MlflowParam(graphene.ObjectType):
+    key = graphene.String()
+    value = graphene.String()
+
+
+class MlflowMetric(graphene.ObjectType):
+    key = graphene.String()
+    value = graphene.Float()
+    timestamp = LongString()
+    step = LongString()
+    dataset_name = graphene.String()
+    dataset_digest = graphene.String()
+    model_id = graphene.String()
+    run_id = graphene.String()
+
+
+class MlflowRunData(graphene.ObjectType):
+    metrics = graphene.List(graphene.NonNull(MlflowMetric))
+    params = graphene.List(graphene.NonNull(MlflowParam))
+    tags = graphene.List(graphene.NonNull(MlflowRunTag))
+
+
+class MlflowRunInfo(graphene.ObjectType):
+    run_id = graphene.String()
+    run_uuid = graphene.String()
+    run_name = graphene.String()
+    experiment_id = graphene.String()
+    user_id = graphene.String()
+    status = graphene.Field(MlflowRunStatus)
+    start_time = LongString()
+    end_time = LongString()
+    artifact_uri = graphene.String()
+    lifecycle_stage = graphene.String()
+
+
+class MlflowRun(graphene.ObjectType):
+    info = graphene.Field(MlflowRunInfo)
+    data = graphene.Field(MlflowRunData)
+    inputs = graphene.Field(MlflowRunInputs)
+    outputs = graphene.Field(MlflowRunOutputs)
+
+
+class MlflowSearchRunsResponse(graphene.ObjectType):
+    runs = graphene.List(graphene.NonNull('mlflow.server.graphql.graphql_schema_extensions.MlflowRunExtension'))
+    next_page_token = graphene.String()
+    apiError = graphene.Field(ApiError)
+
+
+class MlflowGetRunResponse(graphene.ObjectType):
+    run = graphene.Field('mlflow.server.graphql.graphql_schema_extensions.MlflowRunExtension')
+    apiError = graphene.Field(ApiError)
+
+
+class MlflowExperimentTag(graphene.ObjectType):
+    key = graphene.String()
+    value = graphene.String()
+
+
+class MlflowExperiment(graphene.ObjectType):
+    experiment_id = graphene.String()
+    name = graphene.String()
+    artifact_location = graphene.String()
+    lifecycle_stage = graphene.String()
+    last_update_time = LongString()
+    creation_time = LongString()
+    tags = graphene.List(graphene.NonNull(MlflowExperimentTag))
+
+
+class MlflowGetExperimentResponse(graphene.ObjectType):
+    experiment = graphene.Field(MlflowExperiment)
+    apiError = graphene.Field(ApiError)
+
+
+class MlflowSearchModelVersionsInput(graphene.InputObjectType):
+    filter = graphene.String()
+    max_results = LongString()
+    order_by = graphene.List(graphene.String)
+    page_token = graphene.String()
+
+
+class MlflowSearchDatasetsInput(graphene.InputObjectType):
+    experiment_ids = graphene.List(graphene.String)
+
+
+class MlflowGetMetricHistoryBulkIntervalInput(graphene.InputObjectType):
+    run_ids = graphene.List(graphene.String)
+    metric_key = graphene.String()
+    start_step = graphene.Int()
+    end_step = graphene.Int()
+    max_results = graphene.Int()
+
+
+class MlflowListArtifactsInput(graphene.InputObjectType):
+    run_id = graphene.String()
+    run_uuid = graphene.String()
+    path = graphene.String()
+    page_token = graphene.String()
+
+
+class MlflowSearchRunsInput(graphene.InputObjectType):
+    experiment_ids = graphene.List(graphene.String)
+    filter = graphene.String()
+    run_view_type = graphene.Field(MlflowViewType)
+    max_results = graphene.Int()
+    order_by = graphene.List(graphene.String)
+    page_token = graphene.String()
+
+
+class MlflowGetRunInput(graphene.InputObjectType):
+    run_id = graphene.String()
+    run_uuid = graphene.String()
+
+
+class MlflowGetExperimentInput(graphene.InputObjectType):
+    experiment_id = graphene.String()
+
+
+class QueryType(graphene.ObjectType):
+    mlflow_get_experiment = graphene.Field(MlflowGetExperimentResponse, input=MlflowGetExperimentInput())
+    mlflow_get_metric_history_bulk_interval = graphene.Field(MlflowGetMetricHistoryBulkIntervalResponse, input=MlflowGetMetricHistoryBulkIntervalInput())
+    mlflow_get_run = graphene.Field(MlflowGetRunResponse, input=MlflowGetRunInput())
+    mlflow_list_artifacts = graphene.Field(MlflowListArtifactsResponse, input=MlflowListArtifactsInput())
+    mlflow_search_model_versions = graphene.Field(MlflowSearchModelVersionsResponse, input=MlflowSearchModelVersionsInput())
+
+    def resolve_mlflow_get_experiment(self, info, input):
+        input_dict = vars(input)
+        request_message = mlflow.protos.service_pb2.GetExperiment()
+        parse_dict(input_dict, request_message)
+        return mlflow.server.handlers.get_experiment_impl(request_message)
+
+    def resolve_mlflow_get_metric_history_bulk_interval(self, info, input):
+        input_dict = vars(input)
+        request_message = mlflow.protos.service_pb2.GetMetricHistoryBulkInterval()
+        parse_dict(input_dict, request_message)
+        return mlflow.server.handlers.get_metric_history_bulk_interval_impl(request_message)
+
+    def resolve_mlflow_get_run(self, info, input):
+        input_dict = vars(input)
+        request_message = mlflow.protos.service_pb2.GetRun()
+        parse_dict(input_dict, request_message)
+        return mlflow.server.handlers.get_run_impl(request_message)
+
+    def resolve_mlflow_list_artifacts(self, info, input):
+        input_dict = vars(input)
+        request_message = mlflow.protos.service_pb2.ListArtifacts()
+        parse_dict(input_dict, request_message)
+        return mlflow.server.handlers.list_artifacts_impl(request_message)
+
+    def resolve_mlflow_search_model_versions(self, info, input):
+        input_dict = vars(input)
+        request_message = mlflow.protos.model_registry_pb2.SearchModelVersions()
+        parse_dict(input_dict, request_message)
+        return mlflow.server.handlers.search_model_versions_impl(request_message)
+
+
+class MutationType(graphene.ObjectType):
+    mlflow_search_datasets = graphene.Field(MlflowSearchDatasetsResponse, input=MlflowSearchDatasetsInput())
+    mlflow_search_runs = graphene.Field(MlflowSearchRunsResponse, input=MlflowSearchRunsInput())
+
+    def resolve_mlflow_search_datasets(self, info, input):
+        input_dict = vars(input)
+        request_message = mlflow.protos.service_pb2.SearchDatasets()
+        parse_dict(input_dict, request_message)
+        return mlflow.server.handlers.search_datasets_impl(request_message)
+
+    def resolve_mlflow_search_runs(self, info, input):
+        input_dict = vars(input)
+        request_message = mlflow.protos.service_pb2.SearchRuns()
+        parse_dict(input_dict, request_message)
+        return mlflow.server.handlers.search_runs_impl(request_message)

mlflow/server/graphql/graphql_custom_scalars.py

@@ -0,0 +1,24 @@
+import graphene
+from graphql.language.ast import IntValueNode
+
+
+class LongString(graphene.Scalar):
+    """
+    LongString Scalar type to prevent truncation to max integer in JavaScript.
+    """
+
+    description = "Long converted to string to prevent truncation to max integer in JavaScript"
+
+    @staticmethod
+    def serialize(long):
+        return str(long)
+
+    @staticmethod
+    def parse_literal(node):
+        if isinstance(node, IntValueNode):
+            return int(node.value)
+        return None
+
+    @staticmethod
+    def parse_value(value):
+        return int(value)

mlflow/server/graphql/graphql_errors.py

@@ -0,0 +1,15 @@
+import graphene
+
+
+class ErrorDetail(graphene.ObjectType):
+    # NOTE: This is not an exhaustive list, might need to add more things in the future if needed.
+    field = graphene.String()
+    message = graphene.String()
+
+
+class ApiError(graphene.ObjectType):
+    code = graphene.String()
+    message = graphene.String()
+    help_url = graphene.String()
+    trace_id = graphene.String()
+    error_details = graphene.List(ErrorDetail)