QL: Update extractor

ql/extractor/Cargo.toml

@@ -10,7 +10,9 @@ edition = "2018"
 flate2 = "1.0"
 node-types = { path = "../node-types" }
 tree-sitter = "0.19"
-tree-sitter-ql = { git = "https://github.com/erik-krogh/tree-sitter-ql.git", rev = "343cc5873e20510586ade803659ef8ce153bd603" }
+tree-sitter-ql = { git = "https://github.com/tausbn/tree-sitter-ql.git", rev = "2f3fd6f89ca8b2373740d294950a04b9d8e4cdff" }
+tree-sitter-ql-dbscheme = { git = "https://github.com/tausbn/tree-sitter-ql-dbscheme.git"}
+tree-sitter-ql-yaml = {git = "https://github.com/tausbn/tree-sitter-ql.git", rev = "6fe85436702c92b0245d2b91908a12d6b0341f4c"}
 clap = "2.33"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3.3", features = ["env-filter"] }

ql/extractor/src/main.rs

@@ -85,7 +85,13 @@ fn main() -> std::io::Result<()> {
     let file_list = fs::File::open(file_list)?;
 
     let language = tree_sitter_ql::language();
+    let dbscheme = tree_sitter_ql_dbscheme::language();
+    let yaml = tree_sitter_ql_yaml::language();
     let schema = node_types::read_node_types_str("ql", tree_sitter_ql::NODE_TYPES)?;
+    let dbscheme_schema =
+        node_types::read_node_types_str("dbscheme", tree_sitter_ql_dbscheme::NODE_TYPES)?;
+    let yaml_schema = node_types::read_node_types_str("yaml", tree_sitter_ql_yaml::NODE_TYPES)?;
+
     let lines: std::io::Result<Vec<String>> = std::io::BufReader::new(file_list).lines().collect();
     let lines = lines?;
     lines
@@ -105,15 +111,37 @@ fn main() -> std::io::Result<()> {
             let source = std::fs::read(&path)?;
             let code_ranges = vec![];
             let mut trap_writer = trap::Writer::new();
-            extractor::extract(
-                language,
-                "ql",
-                &schema,
-                &mut trap_writer,
-                &path,
-                &source,
-                &code_ranges,
-            )?;
+            if line.ends_with(".dbscheme") {
+                extractor::extract(
+                    dbscheme,
+                    "dbscheme",
+                    &dbscheme_schema,
+                    &mut trap_writer,
+                    &path,
+                    &source,
+                    &code_ranges,
+                )?
+            } else if line.ends_with(".qlpack.yml") {
+                extractor::extract(
+                    yaml,
+                    "yaml",
+                    &yaml_schema,
+                    &mut trap_writer,
+                    &path,
+                    &source,
+                    &code_ranges,
+                )?
+            } else {
+                extractor::extract(
+                    language,
+                    "ql",
+                    &schema,
+                    &mut trap_writer,
+                    &path,
+                    &source,
+                    &code_ranges,
+                )?
+            }
             std::fs::create_dir_all(&src_archive_file.parent().unwrap())?;
             std::fs::copy(&path, &src_archive_file)?;
             write_trap(&trap_dir, path, &trap_writer, trap_compression)

ql/ql/src/codeql_ql/ast/Ast.qll

@@ -32,6 +32,10 @@ class AstNode extends TAstNode {
       node = toQL(this) and
       result = node.getLocation()
     )
+    or
+    result = toGenerateYaml(this).getLocation()
+    or
+    result = toDbscheme(this).getLocation()
   }
 
   predicate hasLocationInfo(
@@ -349,7 +353,7 @@ class Predicate extends TPredicate, AstNode, PredicateOrBuiltin, Declaration {
  * A relation in the database.
  */
 class Relation extends TDBRelation, AstNode, Declaration {
-  QL::DbTable table;
+  Dbscheme::Table table;
 
   Relation() { this = TDBRelation(table) }
 
@@ -358,9 +362,9 @@ class Relation extends TDBRelation, AstNode, Declaration {
    */
   override string getName() { result = table.getTableName().getChild().getValue() }
 
-  private QL::DbColumn getColumn(int i) {
+  private Dbscheme::Column getColumn(int i) {
     result =
-      rank[i + 1](QL::DbColumn column, int child |
+      rank[i + 1](Dbscheme::Column column, int child |
         table.getChild(child) = column
       |
         column order by child
@@ -373,7 +377,7 @@ class Relation extends TDBRelation, AstNode, Declaration {
   /** Gets the `i`th parameter type */
   string getParameterType(int i) {
     // TODO: This is just using the name of the type, not the actual type. Checkout Type.qll
-    result = this.getColumn(i).getColType().getChild().(QL::Token).getValue()
+    result = this.getColumn(i).getColType().getChild().(Dbscheme::Token).getValue()
   }
 
   /**
@@ -2426,11 +2430,13 @@ module YAML {
   class YAMLNode extends TYamlNode, AstNode {
     /** Holds if the predicate is a root node (has no parent) */
     predicate isRoot() { not exists(this.getParent()) }
+
+    override AstNode getParent() { toGenerateYaml(result) = toGenerateYaml(this).getParent() }
   }
 
   /** A YAML comment. */
   class YamlComment extends TYamlCommemt, YAMLNode {
-    QL::YamlComment yamlcomment;
+    Yaml::Comment yamlcomment;
 
     YamlComment() { this = TYamlCommemt(yamlcomment) }
 
@@ -2442,23 +2448,23 @@ module YAML {
 
   /** A YAML entry. */
   class YamlEntry extends TYamlEntry, YAMLNode {
-    QL::YamlEntry yamle;
+    Yaml::Entry yamle;
 
     YamlEntry() { this = TYamlEntry(yamle) }
 
     /** Gets the key of this YAML entry. */
     YamlKey getKey() {
-      exists(QL::YamlKeyvaluepair pair |
+      exists(Yaml::Keyvaluepair pair |
         pair.getParent() = yamle and
         result = TYamlKey(pair.getKey())
       )
     }
 
-    YamlListItem getListItem() { toQL(result).getParent() = yamle }
+    YamlListItem getListItem() { toGenerateYaml(result).getParent() = yamle }
 
     /** Gets the value of this YAML entry. */
     YAMLValue getValue() {
-      exists(QL::YamlKeyvaluepair pair |
+      exists(Yaml::Keyvaluepair pair |
         pair.getParent() = yamle and
         result = TYamlValue(pair.getValue())
       )
@@ -2472,7 +2478,7 @@ module YAML {
 
   /** A YAML key. */
   class YamlKey extends TYamlKey, YAMLNode {
-    QL::YamlKey yamlkey;
+    Yaml::Key yamlkey;
 
     YamlKey() { this = TYamlKey(yamlkey) }
 
@@ -2480,7 +2486,7 @@ module YAML {
      * Gets the value of this YAML key.
      */
     YAMLValue getValue() {
-      exists(QL::YamlKeyvaluepair pair |
+      exists(Yaml::Keyvaluepair pair |
         pair.getKey() = yamlkey and result = TYamlValue(pair.getValue())
       )
     }
@@ -2489,7 +2495,7 @@ module YAML {
 
     /** Gets the value of this YAML value. */
     string getNamePart(int i) {
-      i = 0 and result = yamlkey.getChild(0).(QL::SimpleId).getValue()
+      i = 0 and result = yamlkey.getChild(0).(Yaml::SimpleId).getValue()
       or
       exists(YamlKey child |
         child = TYamlKey(yamlkey.getChild(1)) and
@@ -2511,7 +2517,7 @@ module YAML {
 
   /** A YAML list item. */
   class YamlListItem extends TYamlListitem, YAMLNode {
-    QL::YamlListitem yamllistitem;
+    Yaml::Listitem yamllistitem;
 
     YamlListItem() { this = TYamlListitem(yamllistitem) }
 
@@ -2528,7 +2534,7 @@ module YAML {
 
   /** A YAML value. */
   class YAMLValue extends TYamlValue, YAMLNode {
-    QL::YamlValue yamlvalue;
+    Yaml::Value yamlvalue;
 
     YAMLValue() { this = TYamlValue(yamlvalue) }
 

ql/ql/src/codeql_ql/ast/internal/AstNodes.qll

@@ -14,7 +14,7 @@ newtype TAstNode =
   TClass(QL::Dataclass dc) or
   TCharPred(QL::Charpred pred) or
   TClassPredicate(QL::MemberPredicate pred) or
-  TDBRelation(QL::DbTable table) or
+  TDBRelation(Dbscheme::Table table) or
   TSelect(QL::Select sel) or
   TModule(QL::Module mod) or
   TNewType(QL::Datatype dt) or
@@ -60,11 +60,11 @@ newtype TAstNode =
   TPredicateExpr(QL::PredicateExpr pe) or
   TAnnotation(QL::Annotation annot) or
   TAnnotationArg(QL::AnnotArg arg) or
-  TYamlCommemt(QL::YamlComment yc) or
-  TYamlEntry(QL::YamlEntry ye) or
-  TYamlKey(QL::YamlKey yk) or
-  TYamlListitem(QL::YamlListitem yli) or
-  TYamlValue(QL::YamlValue yv) or
+  TYamlCommemt(Yaml::Comment yc) or
+  TYamlEntry(Yaml::Entry ye) or
+  TYamlKey(Yaml::Key yk) or
+  TYamlListitem(Yaml::Listitem yli) or
+  TYamlValue(Yaml::Value yv) or
   TBuiltinClassless(string ret, string name, string args) { isBuiltinClassless(ret, name, args) } or
   TBuiltinMember(string qual, string ret, string name, string args) {
     isBuiltinMember(qual, ret, name, args)
@@ -124,7 +124,7 @@ private QL::AstNode toQLExpr(AST::AstNode n) {
   n = TDontCare(result)
 }
 
-private QL::AstNode toGenerateYaml(AST::AstNode n) {
+Yaml::AstNode toGenerateYaml(AST::AstNode n) {
   n = TYamlCommemt(result) or
   n = TYamlEntry(result) or
   n = TYamlKey(result) or
@@ -132,6 +132,8 @@ private QL::AstNode toGenerateYaml(AST::AstNode n) {
   n = TYamlValue(result)
 }
 
+Dbscheme::AstNode toDbscheme(AST::AstNode n) { n = TDBRelation(result) }
+
 /**
  * Gets the underlying TreeSitter entity for a given AST node.
  */
@@ -141,8 +143,6 @@ QL::AstNode toQL(AST::AstNode n) {
   or
   result = toQLFormula(n)
   or
-  result = toGenerateYaml(n)
-  or
   result.(QL::ParExpr).getChild() = toQL(n)
   or
   result =
@@ -171,8 +171,6 @@ QL::AstNode toQL(AST::AstNode n) {
   or
   n = TClassPredicate(result)
   or
-  n = TDBRelation(result)
-  or
   n = TSelect(result)
   or
   n = TModule(result)
@@ -233,7 +231,6 @@ module AstConsistency {
     not exists(node.getParent()) and
     not node.getLocation().getStartColumn() = 1 and // startcolumn = 1 <=> top level in file <=> fine to have no parent
     exists(node.toString()) and // <- there are a few parse errors in "global-data-flow-java-1.ql", this way we filter them out.
-    not node instanceof YAML::YAMLNode and // parents for YAML doens't work
     not (node instanceof QLDoc and node.getLocation().getFile().getExtension() = "dbscheme") // qldoc in dbschemes are not hooked up
   }