mirror of
https://github.com/dbt-labs/dbt-core
synced 2025-12-17 19:31:34 +00:00
98 lines
2.9 KiB
Python
Executable File
98 lines
2.9 KiB
Python
Executable File
#!/usr/bin/env python
|
|
import json
|
|
from argparse import ArgumentParser
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Type
|
|
|
|
from dbt.artifacts.schemas.base import VersionedSchema
|
|
from dbt.artifacts.schemas.catalog import CatalogArtifact
|
|
from dbt.artifacts.schemas.freshness import FreshnessExecutionResultArtifact
|
|
from dbt.artifacts.schemas.run import RunResultsArtifact
|
|
from dbt.contracts.graph.manifest import WritableManifest
|
|
from dbt_common.clients.system import write_file
|
|
|
|
|
|
@dataclass
|
|
class ArtifactInfo:
|
|
path: str
|
|
name: str
|
|
json_schema: Dict[str, Any]
|
|
|
|
@classmethod
|
|
def from_artifact_cls(
|
|
cls,
|
|
artifact_cls: Type[VersionedSchema],
|
|
) -> "ArtifactInfo":
|
|
return cls(
|
|
path=artifact_cls.dbt_schema_version.path,
|
|
name=artifact_cls.dbt_schema_version.name,
|
|
json_schema=artifact_cls.json_schema(),
|
|
)
|
|
|
|
def write_schema(self, dest_dir: Path):
|
|
write_file(str(dest_dir / self.path), json.dumps(self.json_schema, indent=2))
|
|
|
|
|
|
@dataclass
|
|
class Arguments:
|
|
artifact: str
|
|
path: Path
|
|
|
|
@classmethod
|
|
def parse(cls) -> "Arguments":
|
|
parser = ArgumentParser(prog="Collect and write dbt arfifact schema")
|
|
parser.add_argument(
|
|
"--path",
|
|
type=Path,
|
|
help="The dir to write artifact schema",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--artifact",
|
|
type=str,
|
|
choices=["manifest", "sources", "run-results", "catalog"],
|
|
help="The name of the artifact to update",
|
|
)
|
|
|
|
parsed = parser.parse_args()
|
|
return cls(artifact=parsed.artifact, path=parsed.path)
|
|
|
|
|
|
def collect_artifact_schema(args: Arguments):
|
|
artifacts = [
|
|
FreshnessExecutionResultArtifact,
|
|
RunResultsArtifact,
|
|
CatalogArtifact,
|
|
# WritableManifest introduces new definitions in hologram which are likely
|
|
# getting persisted across invocations of json_schema and making their
|
|
# way to other written artifacts - so write it last as a short-term fix.
|
|
# https://github.com/dbt-labs/dbt-core/issues/7604
|
|
WritableManifest,
|
|
]
|
|
filtered_artifacts = filter(
|
|
lambda a: a.dbt_schema_version.name == args.artifact or args.artifact is None, artifacts
|
|
)
|
|
artifact_infos = []
|
|
for artifact_cls in filtered_artifacts:
|
|
artifact_infos.append(ArtifactInfo.from_artifact_cls(artifact_cls))
|
|
|
|
if args and args.path is not None:
|
|
for artifact_info in artifact_infos:
|
|
dest_dir = args.path.resolve()
|
|
artifact_info.write_schema(dest_dir)
|
|
else:
|
|
artifacts_dict = {
|
|
artifact_info.name: artifact_info.json_schema for artifact_info in artifact_infos
|
|
}
|
|
print(json.dumps(artifacts_dict))
|
|
|
|
|
|
def main():
|
|
parsed = Arguments.parse()
|
|
collect_artifact_schema(parsed)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|