generator.spoke

  1"""Generate directories and models for new namespaces."""
  2
  3import logging
  4import os
  5import shutil
  6from collections import defaultdict
  7from pathlib import Path
  8from typing import Dict, List, TypedDict
  9
 10import click
 11import lkml
 12import looker_sdk
 13import yaml
 14
 15from .lookml import ViewDict
 16
 17MODEL_SETS_BY_INSTANCE: Dict[str, List[str]] = {
 18    "https://mozilladev.cloud.looker.com": ["mozilla_confidential"],
 19    "https://mozillastaging.cloud.looker.com": ["mozilla_confidential"],
 20    "https://mozilla.cloud.looker.com": ["mozilla_confidential"],
 21}
 22
 23DEFAULT_DB_CONNECTION = "telemetry"
 24
 25
 26class ExploreDict(TypedDict):
 27    """Represent an explore definition."""
 28
 29    type: str
 30    views: List[Dict[str, str]]
 31
 32
 33class NamespaceDict(TypedDict):
 34    """Represent a Namespace definition."""
 35
 36    views: ViewDict
 37    explores: ExploreDict
 38    pretty_name: str
 39    glean_app: bool
 40    connection: str
 41    spoke: str
 42
 43
 44def setup_env_with_looker_creds() -> bool:
 45    """
 46    Set up env with looker credentials.
 47
 48    Returns TRUE if the config is complete.
 49    """
 50    client_id = os.environ.get("LOOKER_API_CLIENT_ID")
 51    client_secret = os.environ.get("LOOKER_API_CLIENT_SECRET")
 52    instance = os.environ.get("LOOKER_INSTANCE_URI")
 53
 54    if client_id is None or client_secret is None or instance is None:
 55        return False
 56
 57    os.environ["LOOKERSDK_BASE_URL"] = instance
 58    os.environ["LOOKERSDK_API_VERSION"] = "4.0"
 59    os.environ["LOOKERSDK_VERIFY_SSL"] = "true"
 60    os.environ["LOOKERSDK_TIMEOUT"] = "120"
 61    os.environ["LOOKERSDK_CLIENT_ID"] = client_id
 62    os.environ["LOOKERSDK_CLIENT_SECRET"] = client_secret
 63
 64    return True
 65
 66
 67def generate_model(
 68    spoke_path: Path, name: str, namespace_defn: NamespaceDict, db_connection: str
 69) -> Path:
 70    """
 71    Generate a model file for a namespace.
 72
 73    We want these to have a nice label and a unique name.
 74    We only import explores and dashboards, as we want those
 75    to auto-import upon generation.
 76
 77    Views are not imported by default, since they should
 78    be added one-by-one if they are included in an explore.
 79    """
 80    logging.info(f"Generating model {name}...")
 81    model_defn = {
 82        "connection": db_connection,
 83        "label": namespace_defn["pretty_name"],
 84    }
 85
 86    # automatically import generated explores for new glean apps
 87    has_explores = len(namespace_defn.get("explores", {})) > 0
 88
 89    path = spoke_path / name / f"{name}.model.lkml"
 90    # lkml.dump may return None, in which case write an empty file
 91    footer_text = f"""
 92# Include files from looker-hub or spoke-default below. For example:
 93{'' if has_explores else '# '}include: "//looker-hub/{name}/explores/*"
 94# include: "//looker-hub/{name}/dashboards/*"
 95# include: "views/*"
 96# include: "explores/*"
 97# include: "dashboards/*"
 98"""
 99    model_text = lkml.dump(model_defn)
100    if model_text is None:
101        path.write_text("")
102    else:
103        path.write_text(model_text + footer_text)
104
105    return path
106
107
108def configure_model(
109    sdk: looker_sdk.methods40.Looker40SDK,
110    model_name: str,
111    db_connection: str,
112    spoke_project: str,
113):
114    """Configure a Looker model by name."""
115    instance = os.environ["LOOKER_INSTANCE_URI"]
116    logging.info(f"Configuring model {model_name}...")
117
118    try:
119        sdk.lookml_model(model_name)
120        logging.info("Model is configured!")
121        return
122    except looker_sdk.error.SDKError:
123        pass
124
125    sdk.create_lookml_model(
126        looker_sdk.models40.WriteLookmlModel(
127            allowed_db_connection_names=[db_connection],
128            name=model_name,
129            project_name=spoke_project,
130        )
131    )
132
133    for model_set_name in MODEL_SETS_BY_INSTANCE[instance]:
134        model_sets = sdk.search_model_sets(name=model_set_name)
135        if len(model_sets) != 1:
136            raise click.ClickException("Error: Found more than one matching model set")
137
138        model_set = model_sets[0]
139        models, _id = model_set.models, model_set.id
140        if models is None or _id is None:
141            raise click.ClickException("Error: Missing models or name from model_set")
142
143        sdk.update_model_set(
144            _id, looker_sdk.models40.WriteModelSet(models=list(models) + [model_name])
145        )
146
147
148def generate_directories(
149    namespaces: Dict[str, NamespaceDict], base_dir: Path, sdk_setup=False
150):
151    """Generate directories and model for a namespace, if it doesn't exist."""
152    seen_spoke_namespaces = defaultdict(list)
153    for namespace, defn in namespaces.items():
154        spoke = defn["spoke"]
155        seen_spoke_namespaces[spoke].append(namespace)
156
157        spoke_dir = base_dir / spoke
158        spoke_dir.mkdir(parents=True, exist_ok=True)
159        print(f"Writing {namespace} to {spoke_dir}")
160        existing_dirs = {p.name for p in spoke_dir.iterdir()}
161
162        if namespace in existing_dirs:
163            continue
164
165        (spoke_dir / namespace).mkdir()
166        for dirname in ("views", "explores", "dashboards"):
167            (spoke_dir / namespace / dirname).mkdir()
168            (spoke_dir / namespace / dirname / ".gitkeep").touch()
169
170        db_connection: str = defn.get("connection", DEFAULT_DB_CONNECTION)
171        generate_model(spoke_dir, namespace, defn, db_connection)
172
173        if sdk_setup:
174            spoke_project = spoke.lstrip("looker-")
175            sdk = looker_sdk.init40()
176            logging.info("Looker SDK 4.0 initialized successfully.")
177            configure_model(sdk, namespace, db_connection, spoke_project)
178
179    # remove directories for namespaces that got removed
180    for spoke in seen_spoke_namespaces.keys():
181        spoke_dir = base_dir / spoke
182        existing_dirs = {p.name for p in spoke_dir.iterdir()}
183
184        for existing_dir in existing_dirs:
185            # make sure the directory belongs to a namespace by checking if a model file exists
186            if (spoke_dir / existing_dir / f"{existing_dir}.model.lkml").is_file():
187                if existing_dir not in seen_spoke_namespaces[spoke]:
188                    # namespace does not exists anymore, remove directory
189                    print(f"Removing {existing_dir} from {spoke_dir}")
190                    shutil.rmtree(spoke_dir / existing_dir)
191
192
193@click.command(help=__doc__)
194@click.option(
195    "--namespaces",
196    default="namespaces.yaml",
197    type=click.File(),
198    help="Path to the namespaces.yaml file.",
199)
200@click.option(
201    "--spoke-dir",
202    default=".",
203    type=click.Path(file_okay=False, dir_okay=True, writable=True),
204    help="Directory containing the Looker spoke.",
205)
206def update_spoke(namespaces, spoke_dir):
207    """Generate updates to spoke project."""
208    _namespaces = yaml.safe_load(namespaces)
209    sdk_setup = setup_env_with_looker_creds()
210    generate_directories(_namespaces, Path(spoke_dir), sdk_setup)