Merge pull request #1014 from atlanhq/DVX-760

Asset import improvements

Author: cmgrote

Diff for: samples/packages/asset-import/src/main/kotlin/com/atlan/pkg/aim/AssetImporter.kt

@@ -251,9 +251,11 @@ class AssetImporter(
 
     /** {@inheritDoc} */
     override fun import(columnsToSkip: Set<String>): ImportResults? {
+        val colsToSkip = columnsToSkip.toMutableSet()
+        colsToSkip.add(Asset.GUID.atlanFieldName)
         if (updateOnly) {
             // If we're only updating, process as before (in-parallel, any order)
-            return super.import(columnsToSkip)
+            return super.import(colsToSkip)
         } else {
             // Otherwise, we need to do multi-pass loading:
             //  - Import assets in tiered order, top-to-bottom
@@ -271,7 +273,7 @@ class AssetImporter(
             typeLoadingOrder.forEach {
                 typeToProcess = it
                 logger.info { "--- Importing $typeToProcess assets... ---" }
-                val results = super.import(columnsToSkip)
+                val results = super.import(colsToSkip)
                 individualResults.add(results)
             }
             return ImportResults.combineAll(Atlan.getDefaultClient(), true, *individualResults.toTypedArray())

Diff for: sdk/src/main/java/com/atlan/util/AssetBatch.java

@@ -20,6 +20,7 @@
 import com.atlan.model.enums.AssetCreationHandling;
 import com.atlan.model.relations.Reference;
 import com.atlan.model.search.FluentSearch;
+import com.atlan.model.search.IndexSearchDSL;
 import com.atlan.serde.Serde;
 import java.io.Closeable;
 import java.io.IOException;
@@ -485,11 +486,12 @@ public synchronized AssetMutationResponse flush() throws AtlanException {
                 } else {
                     builder = client.assets.select(true).where(Asset.QUALIFIED_NAME.in(qualifiedNames));
                 }
-                builder.pageSize(maxSize).stream().forEach(asset -> {
-                    AssetIdentity assetId =
-                            new AssetIdentity(asset.getTypeName(), asset.getQualifiedName(), caseInsensitive);
-                    found.put(assetId, asset.getQualifiedName());
-                });
+                builder.pageSize(Math.max(maxSize * 2, IndexSearchDSL.DEFAULT_PAGE_SIZE)).stream()
+                        .forEach(asset -> {
+                            AssetIdentity assetId =
+                                    new AssetIdentity(asset.getTypeName(), asset.getQualifiedName(), caseInsensitive);
+                            found.put(assetId, asset.getQualifiedName());
+                        });
                 revised = new ArrayList<>();
                 for (Asset asset : _batch) {
                     AssetIdentity assetId =
@@ -519,7 +521,7 @@ public synchronized AssetMutationResponse flush() throws AtlanException {
                             addPartialAsset(asset, revised);
                         } else if (creationHandling == AssetCreationHandling.FULL) {
                             // Still create it (full), if not found and full asset creation is allowed
-                            revised.addAll(_batch);
+                            revised.add(asset);
                         } else {
                             // Otherwise, if it still does not match any fallback and cannot be created, skip it
                             track(skipped, asset);