[sc-27549] Switch to using DiscoveryAPI for dbt.cloud

.github/workflows/ci.yml

@@ -50,10 +50,9 @@ jobs:
           poetry run mypy . --explicit-package-bases
           poetry run bandit -r . -c pyproject.toml
 
-      # TODO(SC-14236): Include __init__.py back to coverage after fixing async testing issues
       - name: Test
         run: |
-          poetry run coverage run --source=metaphor --omit='**/__init__.py' -m pytest
+          poetry run coverage run -m pytest
           poetry run coverage xml
 
       - name: Codecov

.pre-commit-config.yaml

@@ -44,7 +44,8 @@ repos:
     rev: 1.7.8
     hooks:
       - id: bandit
-        args: ['--skip=B101,B106,B404,B603,B607,B608']
+        args: [-c, pyproject.toml]
+        additional_dependencies: ['bandit[toml]']
 
   - repo: https://github.com/pycqa/flake8
     rev: 7.0.0

metaphor/common/entity_id.py

@@ -121,3 +121,20 @@ def dataset_normalized_name(
     return normalize_full_dataset_name(
         ".".join([part for part in [db, schema, table] if part is not None])
     )
+
+
+def parts_to_dataset_entity_id(
+    platform: DataPlatform,
+    account: Optional[str],
+    database: Optional[str] = None,
+    schema: Optional[str] = None,
+    table: Optional[str] = None,
+) -> EntityId:
+    """
+    converts parts of a dataset, its platform and account into a dataset entity ID
+    """
+    return to_dataset_entity_id(
+        dataset_normalized_name(database, schema, table),
+        platform,
+        account,
+    )

metaphor/dbt/cloud/client.py

@@ -15,6 +15,7 @@ class DbtRun(NamedTuple):
     project_id: int
     job_id: int
     run_id: int
+    environment_id: int
 
     def __str__(self) -> str:
         return f"ID = {self.run_id}, project ID = {self.project_id}, job ID = {self.job_id}"
@@ -115,6 +116,7 @@ def get_last_successful_run(
                         project_id=run.get("project_id"),
                         job_id=run.get("job_definition_id"),
                         run_id=run.get("id"),
+                        environment_id=run.get("environment_id"),
                     )
 
             offset += page_size

metaphor/dbt/cloud/discovery_api/__init__.py

@@ -0,0 +1,5 @@
+from .generated.client import Client as DiscoveryAPIClient
+
+__all__ = [
+    "DiscoveryAPIClient",
+]

metaphor/dbt/cloud/discovery_api/apollo-codegen-config.json

@@ -0,0 +1,83 @@
+{
+  "schemaNamespace": "MySchema",
+  "schemaDownload": {
+    "downloadMethod": {
+      "introspection": {
+        "endpointURL": "https://metadata.cloud.getdbt.com/graphql",
+        "httpMethod": {
+          "POST": {}
+        },
+        "includeDeprecatedInputValues": false,
+        "outputFormat": "SDL"
+      }
+    },
+    "downloadTimeout": 60,
+    "headers": [],
+    "outputPath": "./schema.graphql"
+  },
+  "experimentalFeatures": {
+    "clientControlledNullability": true,
+    "legacySafelistingCompatibleOperations": true
+  },
+  "operationManifest": {
+    "generateManifestOnCodeGeneration": false,
+    "path": "/operation/identifiers/path",
+    "version": "persistedQueries"
+  },
+  "input": {
+    "operationSearchPaths": [
+      "/search/path/**/*.graphql"
+    ],
+    "schemaSearchPaths": [
+      "/path/to/schema.graphqls"
+    ]
+  },
+  "output": {
+    "operations": {
+      "absolute": {
+        "accessModifier": "internal",
+        "path": "/absolute/path"
+      }
+    },
+    "schemaTypes": {
+      "moduleType": {
+        "embeddedInTarget": {
+          "accessModifier": "public",
+          "name": "SomeTarget"
+        }
+      },
+      "path": "/output/path"
+    },
+    "testMocks": {
+      "swiftPackage": {
+        "targetName": "SchemaTestMocks"
+      }
+    }
+  },
+  "options": {
+    "additionalInflectionRules": [
+      {
+        "pluralization": {
+          "replacementRegex": "animals",
+          "singularRegex": "animal"
+        }
+      }
+    ],
+    "cocoapodsCompatibleImportStatements": true,
+    "conversionStrategies": {
+      "enumCases": "none",
+      "fieldAccessors": "camelCase",
+      "inputObjects": "camelCase"
+    },
+    "deprecatedEnumCases": "exclude",
+    "operationDocumentFormat": [
+      "definition"
+    ],
+    "pruneGeneratedFiles": false,
+    "schemaDocumentation": "exclude",
+    "selectionSetInitializers": {
+      "localCacheMutations": true
+    },
+    "warningsOnDeprecatedUsage": "exclude"
+  }
+}

metaphor/dbt/cloud/discovery_api/ariadne-codegen.toml

@@ -0,0 +1,5 @@
+[tool.ariadne-codegen]
+schema_path = "schema.graphql"
+queries_path = "queries.graphql"
+async_client = false
+target_package_name = "generated"

metaphor/dbt/cloud/discovery_api/codegen.md

@@ -0,0 +1,38 @@
+# Generate GraphQL client code
+
+## Requirements
+
+- Python >= 3.9
+- `ariadne-codegen`
+
+## Usage
+
+```bash
+cd metaphor/dbt/cloud/discovery_api
+./codegen.sh
+```
+
+## Existing files
+
+### `codegen.sh`
+
+Run this script to get the schema from DBT's Apollo server, and generate the corresponding GraphQL client code.
+
+### `queries.graphql`
+
+The queries we will execute from the extractor class.
+
+### `apollo-codegen-config.json`
+
+Copied from [Full Codegen Configuration Example](https://www.apollographql.com/docs/ios/code-generation/codegen-configuration/#full-codegen-configuration-example) on Apollo's site. The only modifications are:
+
+- `endpointURL`
+- `outputPath`
+
+### `ariadne-codegen.toml`
+
+Controls the behavior of `ariadne-codegen`.
+
+### `schema.graphql`
+
+The upstream DBT GraphQL schema. This file will be downloaded from upstream whenever `codegen.sh` is run.

metaphor/dbt/cloud/discovery_api/codegen.sh

@@ -0,0 +1,17 @@
+#!/usr/bin/env bash
+
+# The tool is called `apollo-ios-cli`: https://www.apollographql.com/docs/ios/code-generation/codegen-cli/
+# It does not mean it's iOS only.
+APOLLO_IOS_CLI_VERSION=1.14.0
+
+wget -c \
+    "https://github.com/apollographql/apollo-ios/releases/download/${APOLLO_IOS_CLI_VERSION}/apollo-ios-cli.tar.gz" -O - | \
+    tar -xz
+
+./apollo-ios-cli fetch-schema --path ./apollo-codegen-config.json
+
+rm -f ./apollo-ios-cli
+
+poetry run ariadne-codegen --config ariadne-codegen.toml
+poetry run black .
+poetry run isort .