Merge pull request #22798 from github/nickrolfe/mad-generator-json

MaD: update generator scripts to support writing JSON data extensions

Author: nickrolfe

misc/scripts/models-as-data/bulk_generate_mad.py

@@ -241,6 +241,7 @@ def generate_models(config, args, project: Project, database_dir: str) -> None:
     generator.with_summaries = should_generate_summaries(project)
     generator.threads = args.codeql_threads
     generator.ram = args.codeql_ram
+    generator.extension_format = args.extension_format
     if config.get("single-file", False):
         generator.single_file = name
     else:
@@ -430,10 +431,11 @@ def download_and_decompress(analyzed_database: dict) -> str:
 def clean_up_mad_destination_for_project(config, name: str):
     target = pathlib.Path(config["destination"], name)
     if config.get("single-file", False):
-        target = target.with_suffix(".model.yml")
-        if target.exists():
-            print(f"Deleting existing MaD file at {target}")
-            target.unlink()
+        for extension in (".model.json", ".model.yml"):
+            model_target = target.with_suffix(extension)
+            if model_target.exists():
+                print(f"Deleting existing MaD file at {model_target}")
+                model_target.unlink()
     elif target.exists():
         print(f"Deleting existing MaD directory at {target}")
         shutil.rmtree(target, ignore_errors=True)
@@ -543,6 +545,12 @@ def main(config, args) -> None:
         help="What `--threads` value to pass to `codeql` (default %(default)s)",
         default=0,
     )
+    parser.add_argument(
+        "--extension-format",
+        choices=["json", "yaml"],
+        default="yaml",
+        help="Format for the generated data extension files (default %(default)s)",
+    )
     args = parser.parse_args()
 
     # Load config file

misc/scripts/models-as-data/convert_extensions.py

@@ -8,25 +8,6 @@
 import tempfile
 
 
-def quote_if_needed(v):
-    # string columns
-    if type(v) is str:
-        return '"' + v + '"'
-    # bool column
-    return str(v)
-
-
-def parseData(data):
-    rows = [{}, {}]
-    for row in data:
-        d = map(quote_if_needed, row)
-        provenance = row[-1]
-        targetRows = rows[1] if provenance.endswith("generated") else rows[0]
-        helpers.insert_update(targetRows, row[0], "      - [" + ", ".join(d) + "]\n")
-
-    return rows
-
-
 class Converter:
     def __init__(self, language, dbDir):
         self.language = language
@@ -38,6 +19,7 @@ def __init__(self, language, dbDir):
         )
         self.extDir = os.path.join(self.codeQlRoot, f"{self.language}/ql/lib/ext/")
         self.dirname = "modelconverter"
+        # N.B. change this to .model.json to change the output format:
         self.modelFileExtension = ".model.yml"
         self.workDir = tempfile.mkdtemp()
 
@@ -63,48 +45,38 @@ def runQuery(self, query):
         )
         return helpers.readData(self.workDir, resultBqrs)
 
-    def asAddsTo(self, rows, predicate):
-        extensions = [{}, {}]
-        for i in range(2):
-            for key in rows[i]:
-                extensions[i][key] = helpers.addsToTemplate.format(
-                    f"codeql/{self.language}-all", predicate, rows[i][key]
-                )
-
-        return extensions
-
-    def getAddsTo(self, query, predicate):
+    def merge_query_results(self, query, predicate, mergers):
         data = self.runQuery(query)
-        rows = parseData(data)
-        return self.asAddsTo(rows, predicate)
-
-    def makeContent(self):
-        summaries = self.getAddsTo("ExtractSummaries.ql", helpers.summaryModelPredicate)
-        sources = self.getAddsTo("ExtractSources.ql", helpers.sourceModelPredicate)
-        sinks = self.getAddsTo("ExtractSinks.ql", helpers.sinkModelPredicate)
-        neutrals = self.getAddsTo("ExtractNeutrals.ql", helpers.neutralModelPredicate)
-        return [
-            helpers.merge(sources[0], sinks[0], summaries[0], neutrals[0]),
-            helpers.merge(sources[1], sinks[1], summaries[1], neutrals[1]),
+        for row in data:
+            provenance = row[-1]
+            namespace = row[0]
+            target_merger = (
+                mergers[1] if provenance.endswith("generated") else mergers[0]
+            )
+            target_merger.add_row(namespace, predicate, row)
+
+    def make_extensions(self):
+        mergers = [
+            helpers.ExtensionMerger(f"codeql/{self.language}-all"),
+            helpers.ExtensionMerger(f"codeql/{self.language}-all"),
         ]
+        self.merge_query_results(
+            "ExtractSummaries.ql", helpers.summaryModelPredicate, mergers
+        )
+        self.merge_query_results(
+            "ExtractSources.ql", helpers.sourceModelPredicate, mergers
+        )
+        self.merge_query_results("ExtractSinks.ql", helpers.sinkModelPredicate, mergers)
+        self.merge_query_results(
+            "ExtractNeutrals.ql", helpers.neutralModelPredicate, mergers
+        )
+        return mergers
+
+    def run(self):
+        mergers = self.make_extensions()
 
-    def save(self, extensions):
         # Create directory if it doesn't exist
         os.makedirs(self.extDir, exist_ok=True)
 
-        # Create a file for each namespace and save models.
-        extensionTemplate = """extensions:
-{0}"""
-        for entry in extensions[0]:
-            with open(self.extDir + "/" + entry + self.modelFileExtension, "w") as f:
-                f.write(extensionTemplate.format(extensions[0][entry]))
-
-        for entry in extensions[1]:
-            with open(
-                self.extDir + "/generated/" + entry + self.modelFileExtension, "w"
-            ) as f:
-                f.write(extensionTemplate.format(extensions[1][entry]))
-
-    def run(self):
-        extensions = self.makeContent()
-        self.save(extensions)
+        mergers[0].save(self.extDir, self.modelFileExtension)
+        mergers[1].save(self.extDir + "/generated", self.modelFileExtension)

misc/scripts/models-as-data/data_extension_file.py

@@ -0,0 +1,76 @@
+import json
+
+
+class DataExtensionFile:
+    """Simple class for representing a data extension file and writing it to disk in a
+    pretty-printed YAML or JSON format."""
+
+    def __init__(self, pack):
+        self.pack = pack
+        # predicate -> dict of tuples (used as an ordered set)
+        self.predicates = {}
+
+    def add_rows(self, predicate, rows):
+        """Add the rows, which can be any collection of any collection of string/boolean values."""
+        if predicate not in self.predicates:
+            self.predicates[predicate] = dict()
+        d = self.predicates[predicate]
+        for row in rows:
+            # To be agnostic about the type of collection used for rows, convert each one to a
+            # tuple. Insert the tuple as a key in the dict, associating it with a dummy value.
+            d[tuple(row)] = None
+
+    def yaml_for_value(self, value):
+        # if is boolean:
+        if isinstance(value, bool):
+            # json.dumps would produce lower-case "true" or "false". Even though that's valid YAML,
+            # we have historically used the upper-case variants.
+            return "True" if value else "False"
+        return json.dumps(value)
+
+    def yaml_for_row(self, row):
+        return f"[{', '.join(self.yaml_for_value(value) for value in row)}]"
+
+    def yaml_for_predicate(self, predicate):
+        return f"""  - addsTo:
+      pack: {self.pack}
+      extensible: {predicate}
+    data:
+      - {'\n      - '.join(self.yaml_for_row(row) for row in self.predicates[predicate].keys())}
+"""
+
+    def write_yaml(self, f):
+        # It would be preferable to use a YAML serialization library, but this lets us control the
+        # formatting e.g. to produce one line per tuple.
+        f.write("# THIS FILE IS AN AUTO-GENERATED MODELS AS DATA FILE. DO NOT EDIT.\n")
+        f.write("extensions:\n")
+        for predicate in sorted(self.predicates.keys()):
+            f.write(self.yaml_for_predicate(predicate))
+
+    def json_for_row(self, row):
+        return "[" + ", ".join(json.dumps(value) for value in row) + "]"
+
+    def json_for_predicate(self, predicate):
+        return f"""    {{
+      "addsTo": {{
+        "pack": "{self.pack}",
+        "extensible": "{predicate}"
+      }},
+      "data": [
+        {',\n        '.join(self.json_for_row(row) for row in self.predicates[predicate].keys())}
+      ]
+    }}"""
+
+    def write_json(self, f):
+        # It would be preferable to use a serializer like Python's json.dumps, but it can't
+        # pretty-print with one line per tuple.
+        f.write("// THIS FILE IS AN AUTO-GENERATED MODELS AS DATA FILE. DO NOT EDIT.\n")
+        f.write("{\n")
+        f.write('  "extensions": [\n')
+        for i, predicate in enumerate(self.predicates):
+            if i > 0:
+                f.write(",\n")
+            f.write(self.json_for_predicate(predicate))
+        f.write("\n")
+        f.write("  ]\n")
+        f.write("}")