[jk] Reuse callback and conditional blocks (mage-ai#4562)

* [jk] Support drag and drop callback and conditional blocks

* [jk] Include callback and conditional blocks in block search

* [jk] Do not add upstream blocks when adding to integration pipeline

Author: johnson-mage

mage_ai/api/resources/SearchResultResource.py

@@ -42,9 +42,7 @@ def filter_results(result: Dict) -> bool:
         return False
 
     if block_type in [
-        BlockType.CALLBACK,
         BlockType.CHART,
-        BlockType.CONDITIONAL,
         BlockType.EXTENSION,
     ]:
         return False

mage_ai/frontend/components/FileBrowser/index.tsx

@@ -13,7 +13,11 @@ import { ThemeContext } from 'styled-components';
 import { useMutation } from 'react-query';
 import { createPortal } from 'react-dom';
 
-import BlockType, { BlockRequestPayloadType, BlockTypeEnum } from '@interfaces/BlockType';
+import BlockType, {
+  ADD_ON_BLOCK_TYPES,
+  BlockRequestPayloadType,
+  BlockTypeEnum,
+} from '@interfaces/BlockType';
 import FileHeaderMenu from './FileHeaderMenu';
 import FileType from '@interfaces/FileType';
 import FlyoutMenu from '@oracle/components/FlyoutMenu';
@@ -323,7 +327,10 @@ function FileBrowser({
               require_unique_name: false,
             },
             block => {
-              if (isIntegrationPipeline && dataExporterBlock) {
+              if (isIntegrationPipeline
+                && dataExporterBlock
+                && !ADD_ON_BLOCK_TYPES.includes(block.type)
+              ) {
                 // @ts-ignore
                 updateDestinationBlock({
                   block: {
@@ -411,7 +418,7 @@ function FileBrowser({
   const selectedBlock = useMemo(() => selectedFile && getBlockFromFile(selectedFile), [
     selectedFile,
   ]);
-  const draggingBlock  = useMemo(() => draggingFile && getBlockFromFile(draggingFile), [
+  const draggingBlock  = useMemo(() => draggingFile && getBlockFromFile(draggingFile, null, true), [
     draggingFile,
   ]);
 

mage_ai/frontend/components/FileBrowser/utils.ts

@@ -22,7 +22,6 @@ import FileType, {
 } from '@interfaces/FileType';
 import { groupBy, prependArray, removeAtIndex, sortByKey } from '@utils/array';
 import { getBlockType } from '@components/FileEditor/utils';
-import { dig } from '@utils/hash';
 import { cleanName, singularize } from '@utils/string';
 
 export function getFullPath(
@@ -123,7 +122,7 @@ export function getBlockFromFile(
   }
 
   const blockTypesToInclude = draggableBlockTypesOnly ? DRAGGABLE_BLOCK_TYPES : BLOCK_TYPES;
-  if (blockTypesToInclude.concat(BlockTypeEnum.DBT).includes(blockType) && validBlockFileExtension(fileName)) {
+  if (blockTypesToInclude.includes(blockType) && validBlockFileExtension(fileName)) {
     const idx = fileName.lastIndexOf('.');
     const extension = fileName.slice(idx + 1);
 

mage_ai/frontend/components/FileEditor/index.tsx

@@ -43,9 +43,6 @@ import {
   getBlockType,
 } from './utils';
 import { find } from '@utils/array';
-import { getBlockFromFile } from '../FileBrowser/utils';
-import { getFullPath } from '@utils/files';
-import { getNonPythonBlockFromFile } from '@components/FileBrowser/utils';
 import { isJsonString } from '@utils/string';
 import { errorOrSuccess, onSuccess } from '@api/utils/response';
 import { onlyKeysPresent } from '@utils/hooks/keyboardShortcuts/utils';

mage_ai/frontend/components/FileEditor/utils.ts

@@ -1,5 +1,6 @@
 import * as osPath from 'path';
 import BlockType, {
+  ADD_ON_BLOCK_TYPES,
   ALL_BLOCK_TYPES,
   ALL_BLOCK_TYPES_WITH_SINGULAR_FOLDERS,
   BlockColorEnum,
@@ -27,7 +28,7 @@ export const getBlockType = (path: string[]): BlockTypeEnum => {
       let part2 = part?.toLowerCase();
 
       if (part2 in ALL_BLOCK_TYPES_WITH_SINGULAR_FOLDERS) {
-        value = part2
+        value = part2;
       } else {
         part2 = singularize(part2);
         if (part2 in ALL_BLOCK_TYPES) {
@@ -50,7 +51,8 @@ export const getBlockType = (path: string[]): BlockTypeEnum => {
     const extensionRegex = new RegExp(`${extensions}$`);
     const fileName = path.join(osPath.sep);
 
-    if (fileName.match(extensionRegex)) {
+    if (fileName.match(extensionRegex)
+      && (path.includes(BlockTypeEnum.DBT) || path.includes('dbts'))) {
       return BlockTypeEnum.DBT;
     }
   }
@@ -121,7 +123,7 @@ export function buildAddBlockRequestPayload(
     blockReqPayload.color = BlockColorEnum.TEAL;
   }
 
-  if (isIntegrationPipeline) {
+  if (isIntegrationPipeline && !ADD_ON_BLOCK_TYPES.includes(blockType)) {
     const dataLoaderBlock: BlockType = find(pipeline.blocks, ({ type }) => BlockTypeEnum.DATA_LOADER === type);
     const transformerBlock: BlockType = find(pipeline.blocks, ({ type }) => BlockTypeEnum.TRANSFORMER === type);
     const upstreamBlocks = [];

mage_ai/frontend/interfaces/BlockType.ts

@@ -74,6 +74,11 @@ export const SIDEKICK_BLOCK_TYPES = [
   BlockTypeEnum.EXTENSION,
 ];
 
+export const ADD_ON_BLOCK_TYPES = [
+  BlockTypeEnum.CALLBACK,
+  BlockTypeEnum.CONDITIONAL,
+];
+
 export enum BlockColorEnum {
   BLUE = 'blue',
   GREY = 'grey',
@@ -88,19 +93,23 @@ export const BLOCK_TYPES = [
   BlockTypeEnum.CUSTOM,
   BlockTypeEnum.DATA_EXPORTER,
   BlockTypeEnum.DATA_LOADER,
+  BlockTypeEnum.DBT,
+  BlockTypeEnum.MARKDOWN,
   BlockTypeEnum.SCRATCHPAD,
   BlockTypeEnum.SENSOR,
-  BlockTypeEnum.MARKDOWN,
   BlockTypeEnum.TRANSFORMER,
 ];
 
 export const DRAGGABLE_BLOCK_TYPES = [
+  BlockTypeEnum.CALLBACK,
+  BlockTypeEnum.CONDITIONAL,
   BlockTypeEnum.CUSTOM,
   BlockTypeEnum.DATA_EXPORTER,
   BlockTypeEnum.DATA_LOADER,
+  BlockTypeEnum.DBT,
+  BlockTypeEnum.MARKDOWN,
   BlockTypeEnum.SCRATCHPAD,
   BlockTypeEnum.SENSOR,
-  BlockTypeEnum.MARKDOWN,
   BlockTypeEnum.TRANSFORMER,
 ];
 

mage_ai/frontend/interfaces/FileType.ts

@@ -78,7 +78,7 @@ export const FILE_EXTENSION_PRIORITY = {
   [FileExtensionEnum.JSON]: 6,
   [FileExtensionEnum.CSV]: 7,
   [FileExtensionEnum.SH]: 8,
-}
+};
 
 export const SUPPORTED_EDITABLE_FILE_EXTENSIONS_REGEX =
   new RegExp(SUPPORTED_EDITABLE_FILE_EXTENSIONS.map(ext => `\.${ext}$`).join('|'));