diff --git a/.vscode/launch.json b/.vscode/launch.json
index c982c0c..a686ca2 100644
--- a/.vscode/launch.json
+++ b/.vscode/launch.json
@@ -27,7 +27,7 @@
       "outFiles": [
         "${workspaceFolder}/packages/dbml-vs-code-extension/dist/extension/*.js"
       ],
-      "preLaunchTask": "yarn: build"
+      "preLaunchTask": "npm: build"
     },
     {
       "name": "Preview Prisma Extension",
@@ -39,7 +39,7 @@
       "outFiles": [
         "${workspaceFolder}/packages/prisma-vs-code-extension/dist/extension/*.js"
       ],
-      "preLaunchTask": "yarn: build:prisma"
+      "preLaunchTask": "npm: build:prisma"
     }
   ]
 }
diff --git a/.vscode/tasks.json b/.vscode/tasks.json
index 38b3778..92d89de 100644
--- a/.vscode/tasks.json
+++ b/.vscode/tasks.json
@@ -45,8 +45,6 @@
         "cwd": "${workspaceFolder}/packages/dbml-vs-code-extension"
       }
     },
-
-    // prisma
     {
       "type": "npm",
       "script": "build:prisma",
@@ -58,6 +56,15 @@
       "options": {
         "cwd": "${workspaceFolder}/packages/prisma-vs-code-extension"
       }
+    },
+    {
+      "type": "npm",
+      "script": "build",
+      "path": "packages/dbml-vs-code-extension",
+      "group": "build",
+      "problemMatcher": [],
+      "label": "npm: build - packages/dbml-vs-code-extension",
+      "detail": "npx vite build && yarn run generate:css"
     }
   ]
 }
diff --git a/.yarn/install-state.gz b/.yarn/install-state.gz
new file mode 100644
index 0000000..5f1bdde
Binary files /dev/null and b/.yarn/install-state.gz differ
diff --git a/packages/extension-shared/extension/types/webviewCommand.ts b/packages/extension-shared/extension/types/webviewCommand.ts
index 0716524..bfcc81f 100644
--- a/packages/extension-shared/extension/types/webviewCommand.ts
+++ b/packages/extension-shared/extension/types/webviewCommand.ts
@@ -12,5 +12,6 @@ export interface WebviewPostMessage {
 export interface SetSchemaCommandPayload {
   type: string;
   payload: JSONTableSchema;
+  message?: string;
   key: string;
 }
diff --git a/packages/extension-shared/extension/views/panel.ts b/packages/extension-shared/extension/views/panel.ts
index 57a9797..9e6b432 100644
--- a/packages/extension-shared/extension/views/panel.ts
+++ b/packages/extension-shared/extension/views/panel.ts
@@ -176,7 +176,15 @@ export class MainPanel {
       MainPanel.diagnosticCollection.clear();
     } catch (error) {
       console.error(JSON.stringify(error));
+
       if (error instanceof DiagnosticError) {
+        // eslint-disable-next-line @typescript-eslint/no-floating-promises
+        this.currentPanel?._panel.webview.postMessage({
+          type: "setSchemaErrorMessage",
+          message: `${error.message}\n Line : ${error.location.start.line}:${error.location.start.column}`,
+          key: document.uri.toString(),
+        });
+
         MainPanel.diagnosticCollection.set(document.uri, [
           new Diagnostic(
             new Range(
diff --git a/packages/extension-shared/src/App.tsx b/packages/extension-shared/src/App.tsx
index db86d16..82f0929 100644
--- a/packages/extension-shared/src/App.tsx
+++ b/packages/extension-shared/src/App.tsx
@@ -2,6 +2,7 @@ import DiagramViewer from "json-table-schema-visualizer/src/components/DiagramVi
 import { useCreateTheme } from "json-table-schema-visualizer/src/hooks/theme";
 import ThemeProvider from "json-table-schema-visualizer/src/providers/ThemeProvider";
 import NoSchemaMessage from "json-table-schema-visualizer/src/components/Messages/NoSchemaMessage";
+import ErrorMessage from "json-table-schema-visualizer/src/components/Messages/ErrorMessage";
 import { type Theme } from "json-table-schema-visualizer/src/types/theme";
 import ScrollDirectionProvider from "json-table-schema-visualizer/src/providers/ScrollDirectionProvider";
 
@@ -16,7 +17,11 @@ const App = () => {
   const { setTheme, theme, themeColors } = useCreateTheme(
     window.EXTENSION_DEFAULT_CONFIG?.theme,
   );
-  const { schema, key } = useSchema();
+  const { schema, key, schemaErrorMessage } = useSchema();
+
+  if (schemaErrorMessage !== null && schema === null) {
+    return <ErrorMessage message={schemaErrorMessage} />;
+  }
 
   if (schema === null) {
     return <NoSchemaMessage />;
diff --git a/packages/extension-shared/src/hooks/schema.ts b/packages/extension-shared/src/hooks/schema.ts
index a845457..9ed6639 100644
--- a/packages/extension-shared/src/hooks/schema.ts
+++ b/packages/extension-shared/src/hooks/schema.ts
@@ -9,12 +9,26 @@ import { type SetSchemaCommandPayload } from "../../extension/types/webviewComma
 export const useSchema = (): {
   schema: JSONTableSchema | null;
   key: string | null;
+  schemaErrorMessage: string | null;
 } => {
+  const [schemaErrorMessage, setSchemaErrorMessage] = useState<string | null>(
+    null,
+  );
   const [schema, setSchema] = useState<JSONTableSchema | null>(null);
   const [schemaKey, setSchemaKey] = useState<string | null>(null);
 
   const updater = (e: MessageEvent): void => {
     const message = e.data as SetSchemaCommandPayload;
+
+    if (
+      message.type === "setSchemaErrorMessage" &&
+      typeof message.message === "string"
+    ) {
+      setSchemaErrorMessage(message.message);
+
+      return;
+    }
+
     if (
       !(message.type === "setSchema" && typeof message.payload === "object")
     ) {
@@ -35,6 +49,7 @@ export const useSchema = (): {
     }
 
     setSchema(message.payload);
+    setSchemaErrorMessage(null);
   };
 
   useEffect(() => {
@@ -47,5 +62,5 @@ export const useSchema = (): {
     };
   }, []);
 
-  return { schema, key: schemaKey };
+  return { schema, key: schemaKey, schemaErrorMessage };
 };
diff --git a/packages/json-table-schema-visualizer/src/components/DiagramViewer/DiagramWrapper.tsx b/packages/json-table-schema-visualizer/src/components/DiagramViewer/DiagramWrapper.tsx
index ec78f8e..2208734 100644
--- a/packages/json-table-schema-visualizer/src/components/DiagramViewer/DiagramWrapper.tsx
+++ b/packages/json-table-schema-visualizer/src/components/DiagramViewer/DiagramWrapper.tsx
@@ -208,6 +208,45 @@ const DiagramWrapper = ({ children }: DiagramWrapperProps) => {
     };
   }, []);
 
+  const onDownload = () => {
+    if (stageRef.current == null) return;
+    const stage = stageRef.current;
+
+    // Save current stage state
+    const originalScale = stage.scaleX();
+    const originalPosition = { ...stage.position() };
+
+    // Reset stage to scale 1 and position 0,0 to get actual content bounds
+    stage.scale({ x: 1, y: 1 });
+    stage.position({ x: 0, y: 0 });
+
+    const contentBounds = stage.getClientRect({ relativeTo: stage });
+
+    // Calculate square dimensions (use the larger dimension)
+    const maxDimension = Math.max(contentBounds.width, contentBounds.height);
+
+    // Center the content in the square
+    const offsetX = (maxDimension - contentBounds.width) / 2;
+    const offsetY = (maxDimension - contentBounds.height) / 2;
+
+    const data = stage.toDataURL({
+      x: contentBounds.x - offsetX,
+      y: contentBounds.y - offsetY,
+      width: maxDimension,
+      height: maxDimension,
+      pixelRatio: 2,
+    });
+
+    // Restore original stage state
+    stage.scale({ x: originalScale, y: originalScale });
+    stage.position(originalPosition);
+
+    const link = document.createElement("a");
+    link.href = data;
+    link.download = `diagram-${Date.now()}.png`;
+    link.click();
+  };
+
   return (
     <>
       <Stage
@@ -229,7 +268,7 @@ const DiagramWrapper = ({ children }: DiagramWrapperProps) => {
         </Layer>
       </Stage>
 
-      <Toolbar onFitToView={fitToView} />
+      <Toolbar onFitToView={fitToView} onDownload={onDownload} />
     </>
   );
 };
diff --git a/packages/json-table-schema-visualizer/src/components/Messages/ErrorMessage.tsx b/packages/json-table-schema-visualizer/src/components/Messages/ErrorMessage.tsx
new file mode 100644
index 0000000..71e1a14
--- /dev/null
+++ b/packages/json-table-schema-visualizer/src/components/Messages/ErrorMessage.tsx
@@ -0,0 +1,14 @@
+import MessageWrapper from "./MessageWrapper";
+
+interface ErrorMessageProps {
+  message: string;
+}
+const ErrorMessage = ({ message }: ErrorMessageProps) => {
+  return (
+    <MessageWrapper>
+      <p>{message}</p>
+    </MessageWrapper>
+  );
+};
+
+export default ErrorMessage;
diff --git a/packages/json-table-schema-visualizer/src/components/Toolbar/Export/Export.tsx b/packages/json-table-schema-visualizer/src/components/Toolbar/Export/Export.tsx
new file mode 100644
index 0000000..22b6d8c
--- /dev/null
+++ b/packages/json-table-schema-visualizer/src/components/Toolbar/Export/Export.tsx
@@ -0,0 +1,17 @@
+import { DownloadIcon } from "lucide-react";
+
+import ToolbarButton from "../Button";
+
+interface ExportButtonProps {
+  onDownload: () => void;
+}
+
+const ExportButton = ({ onDownload }: ExportButtonProps) => {
+  return (
+    <ToolbarButton onClick={onDownload} title="Export">
+      <DownloadIcon />
+    </ToolbarButton>
+  );
+};
+
+export default ExportButton;
diff --git a/packages/json-table-schema-visualizer/src/components/Toolbar/Toolbar.tsx b/packages/json-table-schema-visualizer/src/components/Toolbar/Toolbar.tsx
index 0776235..1b6444e 100644
--- a/packages/json-table-schema-visualizer/src/components/Toolbar/Toolbar.tsx
+++ b/packages/json-table-schema-visualizer/src/components/Toolbar/Toolbar.tsx
@@ -4,14 +4,23 @@ import AutoArrangeTableButton from "./AutoArrage/AutoArrangeTables";
 import ThemeToggler from "./ThemeToggler/ThemeToggler";
 import DetailLevelToggle from "./DetailLevelToggle/DetailLevelToggle";
 import FitToViewButton from "./FitToView/FitToView";
+import ExportButton from "./Export/Export";
 
-const Toolbar = ({ onFitToView }: { onFitToView: () => void }) => {
+const Toolbar = ({
+  onFitToView,
+  onDownload,
+}: {
+  onFitToView: () => void;
+  onDownload: () => void;
+}) => {
   return (
     <div className="flex absolute [&_svg]:w-5 [&_svg]:h-5 px-6 py-1 bottom-14 text-sm bg-gray-100 dark:bg-gray-700 shadow-lg rounded-2xl">
       <AutoArrangeTableButton />
       <DetailLevelToggle />
       <FitToViewButton onClick={onFitToView} />
-      <hr className="w-px h-6 mx-4 my-1 bg-gray-300" />
+      <hr className="mx-4 my-1 w-px h-6 bg-gray-300" />
+      <ExportButton onDownload={onDownload} />
+      <hr className="mx-4 my-1 w-px h-6 bg-gray-300" />
       <ThemeToggler />
     </div>
   );