import csv

import io

import json

import pathlib

import re

import shutil

import time

import click

import sys

from rich.console import Console

from rich.progress import Progress, SpinnerColumn, TextColumn

from rich.table import Table

from diffbot import DiffbotError

from . import ontology

from diffbot import resolve_token

from ._common import get_client

class _DqlGroup(click.Group):

_SECTIONS = [

(

"Commands",

["export"],

),

(

"Advanced Commands",

["init", "ontology", "probe"],

),

]

def invoke(self, ctx):

try:

return super().invoke(ctx)

except (FileNotFoundError, KeyError, DiffbotError) as e:

raise click.ClickException(str(e))

def format_commands(self, ctx, formatter):

section_cmds = {label: [] for label, _ in self._SECTIONS}

uncategorized = []

for name in self.list_commands(ctx):

cmd = self.commands.get(name)

if cmd is None or cmd.hidden:

continue

placed = False

for label, names in self._SECTIONS:

if name in names:

section_cmds[label].append((name, cmd))

placed = True

break

if not placed:

uncategorized.append((name, cmd))

for label, _ in self._SECTIONS:

cmds = section_cmds[label]

if not cmds:

continue

with formatter.section(f"{label}"):

formatter.write_dl(

[(name, cmd.get_short_help_str(formatter.width)) for name, cmd in cmds]

)

if uncategorized:

with formatter.section("Other"):

formatter.write_dl(

[(name, cmd.get_short_help_str(formatter.width)) for name, cmd in uncategorized]

)

@click.group("dql", cls=_DqlGroup)

def dql():

"""Construct and run Diffbot Query Language (DQL) queries against the Knowledge Graph."""

@dql.command("init")

@click.option("--max-age", type=int, default=0, help="Skip the ontology re-download if the local copy is younger than N seconds (default 0 = always refresh).")

@click.option("--keep-tmp", is_flag=True, help="Don't clear ~/.diffbot/tmp.")

def init(max_age: int, keep_tmp: bool):

"""Cache the Diffbot ontology and reset the tmp workspace."""

base = pathlib.Path.home() / ".diffbot"

base.mkdir(exist_ok=True)

tmp = base / "tmp"

if not keep_tmp:

if tmp.exists():

shutil.rmtree(tmp)

tmp.mkdir()

ont = ontology.ONTOLOGY_PATH

fresh = ont.exists() and max_age > 0 and (time.time() - ont.stat().st_mtime) < max_age

if fresh:

click.echo(f"ontology: cached ({int(time.time() - ont.stat().st_mtime)}s old)")

else:

get_client().dql_refresh_ontology(ont)

click.echo(f"ontology: refreshed -> {ont}")

ontology._CACHE.pop("data", None)

if resolve_token():

click.echo("credentials: token found")

else:

click.echo(

"credentials: MISSING -- no Diffbot API token found\n"

" Set the DIFFBOT_API_TOKEN environment variable, or\n"

" write 'DIFFBOT_API_TOKEN=YOUR_TOKEN' to ~/.diffbot/credentials",

err=True,

)

raise SystemExit(2)

@dql.group("ontology")

def ontology_group():

"""Navigate the cached Diffbot ontology (types, fields, taxonomies, enums)."""

@ontology_group.command("types")

def ontology_types():

"""List all entity type names."""

for n in ontology.list_types():

click.echo(n)

@ontology_group.command("composites")

def ontology_composites():

"""List all composite type names."""

for n in ontology.list_composites():

click.echo(n)

@ontology_group.command("enums")

def ontology_enums():

"""List all enum type names."""

for n in ontology.list_enums():

click.echo(n)

@ontology_group.command("taxonomies")

def ontology_taxonomies():

"""List all taxonomy names."""

for n in ontology.list_taxonomies():

click.echo(n)

@ontology_group.command("fields")

@click.argument("type_name", metavar="TYPE")

@click.argument("search", required=False, default=None)

@click.option("--include-deprecated", is_flag=True)

def ontology_fields(type_name: str, search, include_deprecated: bool):

"""List fields of an entity type or composite (e.g. db dql ontology fields Organization)."""

fields = ontology.fields_for(type_name)

for name, meta in ontology.filter_fields(fields, search, include_deprecated=include_deprecated):

click.echo(ontology.format_field(name, meta))

@ontology_group.command("taxonomy")

@click.argument("name")

@click.argument("search", required=False, default=None)

def ontology_taxonomy(name: str, search):

"""List values of a taxonomy."""

for v in ontology.taxonomy_values(name, search):

click.echo(v)

@ontology_group.command("enum")

@click.argument("name")

def ontology_enum(name: str):

"""List values of an enum."""

for v in ontology.enum_values(name):

click.echo(v)

@ontology_group.command("search")

@click.argument("term")

def ontology_search(term: str):

"""Generic fallback: search every 'name' field in the ontology by regex."""

for n in ontology.find_named(term):

click.echo(n)

@dql.command("probe")

@click.argument("queries", nargs=-1, required=True)

@click.option("--workers", type=int, default=8, help="Maximum concurrent requests (default 8).")

@click.option("--json", "as_json", is_flag=True, help="Emit results as a JSON array instead of a text table.")

def probe(queries, workers: int, as_json: bool):

"""Run multiple DQL queries in parallel and print hit counts (size=0) for each."""

reqs = [{"query": q, "size": 0} for q in queries]

results = get_client().dql_parallel(reqs, workers=workers)

rows = [

{"query": q, "hits": r.get("hits"), "results": r.get("results")}

for q, r in zip(queries, results)

]

if as_json:

click.echo(json.dumps(rows, indent=2, ensure_ascii=False))

return

width = max(len(str(row["hits"])) for row in rows)

for row in rows:

click.echo(f"{str(row['hits']).rjust(width)} {row['query']}")

_DEFAULT_SPECS: dict[str, str] = {

"Organization": "name,Name;summary,Summary;nbEmployees,Employees;location.city.name,City;location.country.name,Country",

"Person": "name,Name;$.employments[0].employer.name,Organization;$.employments[0].title,Title;location.city.name,City;location.country.name,Country",

"Article": "date.str,Date;title,Title;author,Author;siteName,Site;language,Language",

"Product": "name,Name;brand,Brand;category,Category;offerPrice,Price;summary,Summary",

"*": "name,Name;description,Description",

}

# Extra fields appended to the stdout default spec for file exports (--out).

# Exports can carry more columns than the terminal table; fill these in per type.

# Format matches _DEFAULT_SPECS: "field.path,Display Name;..." (no leading ';').

_EXPORT_EXTRA_SPECS: dict[str, str] = {

"Organization": "$.categories[*].name,Industries;revenue.value,Revenue;revenue.currency,Revenue Currency;foundingDate.str,Founding Date",

"Person": "allNames,All Names;homepageUri,Website;linkedInUri,LinkedIn;$.employments[1:].employer.name,Employment History;",

"Article": "authorUrl,Author URL;$.categories[*].name,Categories;tags.label,Tags;publisherCountry,Publisher Country;text,Text;html,HTML;",

"Product": "",

"*": "",

}

_TYPE_RE = re.compile(r"\btype:([A-Za-z]+)", re.IGNORECASE)

def _type_for_query(query: str) -> str:

"""Return the entity type named in the query, or '*' if none is found/known."""

m = _TYPE_RE.search(query)

if m and m.group(1) in _DEFAULT_SPECS:

return m.group(1)

return "*"

def _spec_for_query(query: str) -> str:

"""Return the default stdout exportspec for the entity type found in query."""

return _DEFAULT_SPECS[_type_for_query(query)]

def _export_spec_for_query(query: str) -> str:

"""Return the default file-export spec: the stdout default plus any per-type extras."""

t = _type_for_query(query)

base = _DEFAULT_SPECS[t]

extra = _EXPORT_EXTRA_SPECS.get(t, "").strip().strip(";")

return f"{base};{extra}" if extra else base

@dql.command("export")

@click.argument("query")

@click.option("--out", default=None, help="Output file path. If omitted, results are printed to stdout as a table.")

@click.option("--format", "fmt", type=click.Choice(["csv", "xls", "xlsx", "json"]), default="csv", help="Response format for --out (default csv).")

@click.option("--spec", "exportspec", default=None, help="exportspec, e.g. 'name,Name;nbEmployees,Employees;location.city.name,City'")

@click.option("--size", type=int, default=10, help="Page size (default 10).")

@click.option("--from", "from_", type=int, default=0, help="Pagination offset.")

@click.option("--filter", "filter_", default=None, help="Response field filter.")

def export(query: str, out, fmt: str, exportspec, size: int, from_: int, filter_):

"""Execute a DQL query and print to stdout or save to a file."""

if out is not None:

resolved_spec = exportspec or _export_spec_for_query(query)

body = get_client().dql(

query,

size=size,

from_=from_,

format=fmt,

exportspec=resolved_spec or None,

filter=filter_,

raw=True,

)

dest = pathlib.Path(out).expanduser()

dest.parent.mkdir(parents=True, exist_ok=True)

dest.write_bytes(body)

click.echo(f"saved: {dest} ({len(body)} bytes)")

return

client = get_client()

resolved_spec = exportspec or _spec_for_query(query)

if sys.stdout.isatty():

with Progress(SpinnerColumn(), TextColumn("[progress.description]{task.description}"), transient=True) as progress:

progress.add_task(description=f"[Querying...] {query}", total=None)

hits = client.dql(query, size=0).get("hits", 0)

csv_bytes = client.dql(query, size=size, from_=from_, format="csv", exportspec=resolved_spec or None, filter=filter_, raw=True)

else:

hits = client.dql(query, size=0).get("hits", 0)

csv_bytes = client.dql(query, size=size, from_=from_, format="csv", exportspec=resolved_spec or None, filter=filter_, raw=True)

click.echo("")

click.echo(f"[Query] {query}")

reader = csv.reader(io.StringIO(csv_bytes.decode("utf-8")))

rows = list(reader)

if len(rows) < 2:

click.echo("(no results)")

return

headers, *data_rows = rows

table = Table(show_header=True, header_style="bold", show_lines=True)

for i, header in enumerate(headers):

if i < 3:

table.add_column(header, overflow="ellipsis", max_width=40, no_wrap=True)

else:

table.add_column(header, overflow="ellipsis", max_width=20, no_wrap=True)

for row in data_rows:

table.add_row(*row)

Console().print(table)

click.echo(f"{len(data_rows)} of {hits} records")

click.echo("")

click.echo("Options:")

click.echo(" --size N return N records (default 10; -1 for all)")

click.echo(" --from N skip the first N records (pagination offset)")

click.echo("")