Merge pull request #80 from JupiterOne/INT-11063

use on-disk dkt; remove findEntity calls
JupiterOne-Archives · May 20, 2024 · e84e53b · e84e53b
2 parents 6006c39 + 22ad148
commit e84e53b
Show file tree

Hide file tree

Showing 8 changed files with 186 additions and 59 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -1,6 +1,7 @@
 FROM node:18-alpine
 
 ENV JUPITERONE_INTEGRATION_DIR=/opt/jupiterone/integration
+ENV USE_ON_DISK_DKT='1'
 ENV NO_COLLECTION_METRICS='1'
 
 COPY package.json yarn.lock tsconfig.json LICENSE ${JUPITERONE_INTEGRATION_DIR}/

diff --git a/package.json b/package.json
@@ -44,6 +44,7 @@
   },
   "dependencies": {
     "@jupiterone/integration-sdk-http-client": "^12.8.1",
+    "lmdb": "^3.0.8",
     "node-fetch": "^2.7.0",
     "node-match-path": "^0.4.4"
   }

diff --git a/src/constants.ts b/src/constants.ts
@@ -70,6 +70,9 @@ export const entities: Record<EntityConstantKeys, StepEntityMetadata> = {
     resourceName: 'ArtifactCodeModule',
     _type: 'artifactory_artifact_codemodule',
     _class: ['CodeModule'],
+    indexMetadata: {
+      enabled: false,
+    },
   },
   BUILD: {
     resourceName: 'Build',

diff --git a/src/steps/access.ts b/src/steps/access.ts
@@ -61,13 +61,15 @@ export async function fetchUsers({
 
     if (user.groups) {
       for (const groupName of user.groups) {
-        const groupEntity = await jobState.findEntity(getGroupKey(groupName));
-        if (groupEntity) {
+        const groupKey = getGroupKey(groupName);
+        if (jobState.hasKey(groupKey)) {
           await jobState.addRelationship(
             createDirectRelationship({
               _class: RelationshipClass.HAS,
-              from: groupEntity,
-              to: userEntity,
+              fromKey: groupKey,
+              fromType: entities.GROUP._type,
+              toKey: userEntity._key,
+              toType: entities.USER._type,
             }),
           );
         }
@@ -160,16 +162,18 @@ export async function fetchAccessTokens({
 
     const [, , username] = token.subject.split('/');
 
-    const userEntity = await jobState.findEntity(getUserKey(username));
+    const userKey = getUserKey(username);
 
-    if (userEntity) {
+    if (jobState.hasKey(userKey)) {
       await Promise.all([
         jobState.addEntity(tokenEntity),
         jobState.addRelationship(
           createDirectRelationship({
             _class: RelationshipClass.ASSIGNED,
-            from: tokenEntity,
-            to: userEntity,
+            fromKey: tokenEntity._key,
+            fromType: entities.ACCESS_TOKEN._type,
+            toKey: userKey,
+            toType: entities.USER._type,
           }),
         ),
       ]);

diff --git a/src/steps/builds.ts b/src/steps/builds.ts
@@ -3,8 +3,6 @@ import {
   createIntegrationEntity,
   IntegrationStep,
   createDirectRelationship,
-  JobState,
-  Entity,
   RelationshipClass,
 } from '@jupiterone/integration-sdk-core';
 
@@ -17,11 +15,8 @@ export function getBuildKey(name: string): string {
   return `artifactory_build:${name}`;
 }
 
-async function createBuildEntity(
-  jobState: JobState,
-  build: ArtifactoryBuild,
-): Promise<Entity> {
-  const buildEntity = createIntegrationEntity({
+function createBuildEntity(build: ArtifactoryBuild) {
+  return createIntegrationEntity({
     entityData: {
       source: build,
       assign: {
@@ -33,8 +28,6 @@ async function createBuildEntity(
       },
     },
   });
-
-  return jobState.addEntity(buildEntity);
 }
 
 export async function fetchBuilds({
@@ -46,21 +39,22 @@ export async function fetchBuilds({
 
   try {
     await apiClient.iterateBuilds(async (build) => {
-      const buildEntity =
-        (await jobState.findEntity(getBuildKey(build.name))) ||
-        (await createBuildEntity(jobState, build));
+      const buildEntity = createBuildEntity(build);
+      if (!jobState.hasKey(buildEntity._key)) {
+        await jobState.addEntity(buildEntity);
+      }
 
       for (const artifactUri of build.artifacts) {
-        const artifactEntity = await jobState.findEntity(
-          getArtifactKey(artifactUri),
-        );
+        const artifactKey = getArtifactKey(artifactUri);
 
-        if (artifactEntity) {
+        if (jobState.hasKey(artifactKey)) {
           await jobState.addRelationship(
             createDirectRelationship({
               _class: RelationshipClass.CREATED,
-              from: buildEntity,
-              to: artifactEntity,
+              fromKey: buildEntity._key,
+              fromType: entities.BUILD._type,
+              toKey: artifactKey,
+              toType: entities.ARTIFACT_CODEMODULE._type,
             }),
           );
         }

diff --git a/src/steps/permissions.ts b/src/steps/permissions.ts
@@ -130,19 +130,19 @@ async function createPermissionAllowsRepositoryGroupRelationships(
       continue;
     }
 
-    const group = await jobState.findEntity(
-      getRepositoryGroupKey(repositoryName),
-    );
+    const groupKey = getRepositoryGroupKey(repositoryName);
 
-    if (group) {
+    if (jobState.hasKey(groupKey)) {
       await jobState.addRelationship(
         createDirectRelationship({
           _class: RelationshipClass.ALLOWS,
-          from: permissionEntity,
-          to: group,
+          fromKey: permissionEntity._key,
+          fromType: entities.PERMISSION._type,
+          toKey: groupKey,
+          toType: entities.REPOSITORY_GROUP._type,
           properties: {
             _type: relationships.PERMISSION_ALLOWS_REPOSITORY_GROUP._type,
-            _key: `${permissionEntity._key}|allows|${group._key}`,
+            _key: `${permissionEntity._key}|allows|${groupKey}`,
           },
         }),
       );
@@ -195,14 +195,16 @@ export async function fetchPermissions({
     for (const [groupName, permissions] of Object.entries(
       constructPermissionsMap(permission, 'groups'),
     )) {
-      const groupEntity = await jobState.findEntity(getGroupKey(groupName));
+      const groupKey = getGroupKey(groupName);
 
-      if (groupEntity) {
+      if (jobState.hasKey(groupKey)) {
         await jobState.addRelationship(
           createDirectRelationship({
             _class: RelationshipClass.ASSIGNED,
-            from: permissionEntity,
-            to: groupEntity,
+            fromKey: permissionEntity._key,
+            fromType: entities.PERMISSION._type,
+            toKey: groupKey,
+            toType: entities.GROUP._type,
             properties: permissions,
           }),
         );
@@ -213,14 +215,16 @@ export async function fetchPermissions({
     for (const [username, permissions] of Object.entries(
       constructPermissionsMap(permission, 'users'),
     )) {
-      const userEntity = await jobState.findEntity(getUserKey(username));
+      const userKey = getUserKey(username);
 
-      if (userEntity) {
+      if (jobState.hasKey(userKey)) {
         await jobState.addRelationship(
           createDirectRelationship({
             _class: RelationshipClass.ASSIGNED,
-            from: permissionEntity,
-            to: userEntity,
+            fromKey: permissionEntity._key,
+            fromType: entities.PERMISSION._type,
+            toKey: userKey,
+            toType: entities.USER._type,
             properties: permissions,
           }),
         );

diff --git a/src/steps/repositories/index.ts b/src/steps/repositories/index.ts
@@ -16,8 +16,7 @@ import {
   Steps,
 } from '../../constants';
 import { IntegrationConfig } from '../../types';
-import { compareEntities } from '../../utils';
-import { createArtifactEntity, getArtifactKey } from './converters';
+import { createArtifactEntity } from './converters';
 
 export function getRepositoryKey(name: string): string {
   return `artifactory_repository:${name}`;
@@ -139,28 +138,18 @@ export async function fetchArtifacts({
         await apiClient.iterateRepositoryArtifacts(
           repoKey,
           async (artifact) => {
-            const artifactKey = getArtifactKey(artifact.uri);
-            let artifactEntity = await jobState.findEntity(artifactKey);
+            const artifactEntity = createArtifactEntity(artifact, packageType);
 
-            if (!artifactEntity) {
-              artifactEntity = createArtifactEntity(artifact, packageType);
+            if (!jobState.hasKey(artifactEntity._key)) {
               await jobState.addEntity(artifactEntity);
-            } else {
-              const duplicatedArtifact = createArtifactEntity(
-                artifact,
-                packageType,
-              );
-              const duplicateEntityReport = compareEntities(
-                artifactEntity,
-                duplicatedArtifact,
-              );
-              logger.debug(duplicateEntityReport, 'Duplicate entity report.');
             }
 
             const repoArtifactRelationship = createDirectRelationship({
               _class: RelationshipClass.HAS,
-              from: repositoryEntity,
-              to: artifactEntity,
+              fromKey: repositoryEntity._key,
+              fromType: entities.REPOSITORY._type,
+              toKey: artifactEntity._key,
+              toType: entities.ARTIFACT_CODEMODULE._type,
             });
 
             if (!jobState.hasKey(repoArtifactRelationship._key)) {