run mypy on dagster-graphql (#7075)

* enable mypy on dagster-graphql

* fix mypy

* fix lint

.buildkite/dagster-buildkite/dagster_buildkite/module_build_spec.py

@@ -7,7 +7,6 @@
 MYPY_EXCLUDES = [
     "python_modules/dagit",
     "python_modules/automation",
-    "python_modules/dagster-graphql",
     "python_modules/libraries/dagster-databricks",
     "python_modules/libraries/dagster-dbt",
     "python_modules/libraries/dagster-docker",

python_modules/dagster-graphql/dagster_graphql/client/client.py

@@ -106,14 +106,9 @@ def _get_repo_locations_and_names_with_pipeline(self, pipeline_name: str) -> Lis
         repo_connection_status = query_res["__typename"]
         if repo_connection_status == "RepositoryConnection":
             valid_nodes: Iterable[PipelineInfo] = chain(
-                map(PipelineInfo.from_node, query_res["nodes"])
+                *map(PipelineInfo.from_node, query_res["nodes"])
             )
-            return [
-                repo_node_tuple
-                for repo_node_tuple_lst in valid_nodes
-                for repo_node_tuple in repo_node_tuple_lst
-                if repo_node_tuple.pipeline_name == pipeline_name
-            ]
+            return [info for info in valid_nodes if info.pipeline_name == pipeline_name]
         else:
             raise DagsterGraphQLClientError(repo_connection_status, query_res["message"])
 

python_modules/dagster-graphql/dagster_graphql/implementation/fetch_assets.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Mapping
+from typing import TYPE_CHECKING, Dict, Mapping
 
 from dagster import AssetKey, DagsterEventType, EventRecordsFilter, check, seven
 from dagster.core.events import ASSET_EVENTS
@@ -64,7 +64,7 @@ def get_asset_nodes_by_asset_key(graphene_info) -> Mapping[AssetKey, "GrapheneAs
 
     from ..schema.asset_graph import GrapheneAssetNode
 
-    asset_nodes_by_asset_key = {}
+    asset_nodes_by_asset_key: Dict[AssetKey, GrapheneAssetNode] = {}
     for location in graphene_info.context.repository_locations:
         for repository in location.get_repositories().values():
             for external_asset_node in repository.get_external_asset_nodes():

python_modules/dagster-graphql/dagster_graphql/implementation/utils.py

@@ -1,10 +1,12 @@
 import sys
 from collections import namedtuple
+from typing import cast
 
 from graphql.execution.base import ResolveInfo
 
 from dagster import check
 from dagster.core.host_representation import GraphSelector, PipelineSelector
+from dagster.core.workspace.context import BaseWorkspaceRequestContext
 from dagster.utils.error import serializable_error_info_from_exc_info
 
 
@@ -23,7 +25,8 @@ def _fn(self, graphene_info, *args, **kwargs):  # pylint: disable=unused-argumen
 def assert_permission(graphene_info: ResolveInfo, permission: str) -> None:
     from dagster_graphql.schema.errors import GrapheneUnauthorizedError
 
-    if not graphene_info.context.has_permission(permission):
+    context = cast(BaseWorkspaceRequestContext, graphene_info.context)
+    if not context.has_permission(permission):
         raise UserFacingGraphQLError(GrapheneUnauthorizedError())
 
 

python_modules/dagster-graphql/dagster_graphql/schema/asset_graph.py

@@ -160,11 +160,10 @@ def resolve_assetMaterializations(
     ) -> List[GrapheneMaterializationEvent]:
         from ..implementation.fetch_assets import get_asset_materializations
 
+        beforeTimestampMillis: Optional[str] = kwargs.get("beforeTimestampMillis")
         try:
             before_timestamp = (
-                int(kwargs.get("beforeTimestampMillis")) / 1000.0
-                if kwargs.get("beforeTimestampMillis")
-                else None
+                int(beforeTimestampMillis) / 1000.0 if beforeTimestampMillis else None
             )
         except ValueError:
             before_timestamp = None

python_modules/dagster-graphql/dagster_graphql_tests/client_tests/test_reload_repository_location.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 import pytest
 from dagster_graphql import DagsterGraphQLClientError, ReloadRepositoryLocationStatus
 
@@ -86,13 +88,12 @@ def test_failure_with_query_error(mock_client: MockClient):
         mock_client.python_client.reload_repository_location("foo")
 
 
-class TestReloadRepositoryLocationWithClient(
-    make_graphql_context_test_suite(
-        context_variants=[
-            GraphQLContextVariant.non_launchable_in_memory_instance_managed_grpc_env()
-        ]
-    )
-):
+BaseTestSuite: Any = make_graphql_context_test_suite(
+    context_variants=[GraphQLContextVariant.non_launchable_in_memory_instance_managed_grpc_env()]
+)
+
+
+class TestReloadRepositoryLocationWithClient(BaseTestSuite):
     def test_reload_location_real(self, graphql_client):
         assert (
             graphql_client.reload_repository_location("test").status

python_modules/dagster-graphql/dagster_graphql_tests/client_tests/test_shutdown_repository_location.py

@@ -1,4 +1,5 @@
 import time
+from typing import Any
 
 from dagster_graphql import ShutdownRepositoryLocationStatus
 
@@ -9,12 +10,12 @@
     make_graphql_context_test_suite,
 )
 
+BaseTestSuite: Any = make_graphql_context_test_suite(
+    context_variants=[GraphQLContextVariant.non_launchable_sqlite_instance_deployed_grpc_env()]
+)
+
 
-class TestShutdownRepositoryLocation(
-    make_graphql_context_test_suite(
-        context_variants=[GraphQLContextVariant.non_launchable_sqlite_instance_deployed_grpc_env()]
-    )
-):
+class TestShutdownRepositoryLocation(BaseTestSuite):
     def test_shutdown_repository_location(self, graphql_client, graphql_context):
         origin = next(iter(graphql_context.get_workspace_snapshot().values())).origin
         origin.create_client().heartbeat()