[SPEC] Allow nested struct fields in SchemaDatasetFacet

Signed-off-by: Martynov Maxim <martinov_m_s_@mail.ru>
OpenLineage · Apr 3, 2024 · db2dabe · db2dabe
1 parent 6bb5142
commit db2dabe
Show file tree

Hide file tree

Showing 3 changed files with 117 additions and 47 deletions.
diff --git a/client/java/generator/src/main/java/io/openlineage/client/JavaPoetGenerator.java b/client/java/generator/src/main/java/io/openlineage/client/JavaPoetGenerator.java
@@ -735,7 +735,12 @@ public TypeName visit(ArrayResolvedType arrayType) {
 
       @Override
       public TypeName visit(TypeResolver.EnumResolvedType enumType) {
-         return ClassName.get(containerClass, enumType.getParentName() + "." + enumType.getName());
+        return ClassName.get(containerClass, enumType.getParentName() + "." + enumType.getName());
+      }
+
+      @Override
+      public TypeName visit(TypeResolver.RefResolvedType refType) {
+        return ClassName.get(containerClass, refType.getName());
       }
     });
   }

diff --git a/client/java/generator/src/main/java/io/openlineage/client/TypeResolver.java b/client/java/generator/src/main/java/io/openlineage/client/TypeResolver.java
@@ -43,7 +43,6 @@
 public class TypeResolver {
 
   private Map<String, ObjectResolvedType> types = new HashMap<>();
-  private Set<String> referencedTypes = new HashSet<>();
   private Set<String> baseTypes = new HashSet<>();
 
   private Map<URL, ResolvedType> rootResolvedTypePerURL = new HashMap<URL, TypeResolver.ResolvedType>();
@@ -75,7 +74,6 @@ public ResolvedType visit(ObjectType objectType) {
           List<Field> properties = objectType.getProperties();
           currentObjectName = currentName;
           List<ResolvedField> resolvedFields = new ArrayList<>(properties.size());
-          resolvedFields.addAll(resolveFields(properties));
           ObjectResolvedType objectResolvedType = new ObjectResolvedType(
               container,
               asList(objectType),
@@ -88,6 +86,7 @@ public ResolvedType visit(ObjectType objectType) {
           if (types.put(key, objectResolvedType) != null) {
             throw new RuntimeException("Duplicated type: " + objectResolvedType.getName());
           };
+          resolvedFields.addAll(resolveFields(properties));
           return objectResolvedType;
         }
 
@@ -104,28 +103,6 @@ private List<ResolvedField> resolveFields(List<Field> properties) {
             currentName = previousCurrentName + currentProperty;
             ResolvedField resolvedField = new ResolvedField(property, visit(property.getType()));
             resolvedFields.add(resolvedField);
-            referencedTypes.add(resolvedField.getType().accept(new ResolvedTypeVisitor<String>() {
-              @Override
-              public String visit(PrimitiveResolvedType primitiveType) {
-                return primitiveType.getName();
-              }
-
-              @Override
-              public String visit(ObjectResolvedType objectType) {
-                return objectType.getName();
-              }
-
-              @Override
-              public String visit(ArrayResolvedType arrayType) {
-                return visit(arrayType.items);
-              }
-
-              @Override
-              public String visit(EnumResolvedType enumType) {
-                return enumType.getName();
-              }
-
-            }));
           }
           currentName = previousCurrentName;
           return resolvedFields;
@@ -164,7 +141,18 @@ public ResolvedType visit(AllOfType allOfType) {
           ResolvedType additionalPropertiesType = null;
           Set<ObjectResolvedType> parents = new LinkedHashSet<>();
           for (Type child : children) {
-            ObjectResolvedType resolvedChildType = (ObjectResolvedType) visit(child);
+            ResolvedType childType = visit(child);
+            ObjectResolvedType resolvedChildType;
+
+            if (childType instanceof RefResolvedType) {
+              RefResolvedType refType = (RefResolvedType) childType;
+              resolvedChildType = types.get(refType.getFullName());
+              if (resolvedChildType == null) {
+                throw new RuntimeException("Unknown type: " + refType.getFullName());
+              }
+            } else {
+              resolvedChildType = (ObjectResolvedType) childType;
+            }
             List<ObjectType> objectTypes = resolvedChildType.getObjectTypes();
             if (!currentName.equals(resolvedChildType.getName())) {
               // base interface
@@ -208,7 +196,7 @@ public ResolvedType visit(RefType refType) {
           }
           String key = refContainer + "." + typeName;
           if (types.containsKey(key)) {
-            return types.get(key);
+            return new RefResolvedType(refContainer, typeName);
           }
           if (anchorIndex > 0) {
             throw new RuntimeException("This ref should have been resolved already: " + refContainer + " " + refType.getPointer() + " => "+ key + " keys: " + types.keySet());
@@ -283,6 +271,8 @@ interface ResolvedTypeVisitor<T> {
 
     T visit(EnumResolvedType enumType);
 
+    T visit(RefResolvedType refType);
+
     default T visit(ResolvedType type) {
       try {
         return type == null ? null : type.accept(this);
@@ -303,6 +293,8 @@ public static class DefaultResolvedTypeVisitor<T> implements ResolvedTypeVisitor
 
     public T visit(EnumResolvedType enumResolvedType) { return null; }
 
+    public T visit(RefResolvedType refType) { return null; }
+
   }
 
   static class PrimitiveResolvedType implements ResolvedType {
@@ -338,6 +330,14 @@ public boolean equals(Object o) {
       return primitiveType.equals(that.primitiveType);
     }
 
+    @Override
+    public String toString() {
+      if (getFormat() == null) {
+        return "PrimitiveType{" + getName() + "}";
+      }
+      return "PrimitiveType{name: " + getName() + ", format: " + getFormat() +  "}";
+    }
+
     @Override
     public int hashCode() {
       return Objects.hash(primitiveType);
@@ -474,6 +474,52 @@ public int hashCode() {
     }
   }
 
+  static class RefResolvedType implements ResolvedType {
+
+    private String container;
+    private String name;
+
+    public RefResolvedType(String container, String name) {
+      this.container = container;
+      this.name = name;
+    }
+
+    public String getFullName() {
+      return container + "." + name;
+    }
+
+    public String getName() {
+      return name;
+    }
+
+    @Override
+    public String toString() {
+      return "RefResolvedType{" + getFullName() + "}";
+    }
+
+    @Override
+    public <T> T accept(ResolvedTypeVisitor<T> visitor) {
+      return visitor.visit(this);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) {
+        return true;
+      }
+      if (o == null || getClass() != o.getClass()) {
+        return false;
+      }
+      RefResolvedType that = (RefResolvedType) o;
+      return container.equals(that.container) && name.equals(that.name);
+    }
+
+    @Override
+    public int hashCode() {
+      return Objects.hash(container, name);
+    }
+  }
+
   static class ArrayResolvedType implements ResolvedType {
 
     private final ArrayType arrayType;
@@ -505,6 +551,11 @@ public boolean equals(Object o) {
       return arrayType.equals(that.arrayType) && items.equals(that.items);
     }
 
+    @Override
+    public String toString() {
+      return "ArrayResolvedType{type: " + arrayType.toString() + ", items: " + items.toString() +  "}";
+    }
+
     @Override
     public int hashCode() {
       return Objects.hash(arrayType, items);

diff --git a/spec/facets/SchemaDatasetFacet.json b/spec/facets/SchemaDatasetFacet.json
@@ -1,7 +1,38 @@
 {
   "$schema": "https://json-schema.org/draft/2020-12/schema",
-  "$id": "https://openlineage.io/spec/facets/1-0-0/SchemaDatasetFacet.json",
+  "$id": "https://openlineage.io/spec/facets/1-0-2/SchemaDatasetFacet.json",
   "$defs": {
+    "SchemaDatasetFacetField": {
+      "type": "object",
+      "properties": {
+        "name": {
+          "description": "The name of the field.",
+          "type": "string",
+          "example": "column1"
+        },
+        "type": {
+          "description": "The type of the field.",
+          "type": "string",
+          "example": "VARCHAR|INT|..."
+        },
+        "description": {
+          "description": "The description of the field.",
+          "type": "string"
+        },
+        "fields": {
+          "description": "Nested struct fields.",
+          "type": "array",
+          "items": {
+            "$ref": "#/$defs/SchemaDatasetFacetField"
+          }
+        },
+        "repeated": {
+          "description": "Is the field repeated (e.g. array).",
+          "type": "boolean"
+        }
+      },
+      "required": ["name"]
+    },
     "SchemaDatasetFacet": {
       "allOf": [
         {
@@ -11,27 +42,10 @@
           "type": "object",
           "properties": {
             "fields": {
-              "description": "The fields of the table.",
+              "description": "The fields of the data source.",
               "type": "array",
               "items": {
-                "type": "object",
-                "properties": {
-                  "name": {
-                    "description": "The name of the field.",
-                    "type": "string",
-                    "example": "column1"
-                  },
-                  "type": {
-                    "description": "The type of the field.",
-                    "type": "string",
-                    "example": "VARCHAR|INT|..."
-                  },
-                  "description": {
-                    "description": "The description of the field.",
-                    "type": "string"
-                  }
-                },
-                "required": ["name"]
+                "$ref": "#/$defs/SchemaDatasetFacetField"
               }
             }
           }