CB-1393 Web resources quote error model. Data transfer quota

samples/sample-databases/cloudbeaver.conf

@@ -38,7 +38,14 @@
         supportsCustomConnections: false,
 
         publicCredentialsSaveEnabled: true,
-        adminCredentialsSaveEnabled: true
+        adminCredentialsSaveEnabled: true,
+
+        resourceQuotas: {
+            dataExportFileSizeLimit: 10000000,
+            sqlMaxRunningQueries: 3,
+            sqlResultSetRowsLimit: 100000,
+            sqlResultSetMemoryLimit: 2000000
+        }
 
         /*
         enabledDrivers: [

server/bundles/io.cloudbeaver.server/schema/service.core.graphqls

@@ -68,6 +68,7 @@ type AsyncTaskInfo {
 type ServerError {
     message: String
     errorCode: String
+    errorType: String
     stackTrace: String
     causedBy: ServerError
 }

server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/DBWebException.java

@@ -43,6 +43,7 @@ public class DBWebException extends DBException implements GraphQLError {
     public static final String ERROR_CODE_LICENSE_DENIED = "licenseRequired";
     public static final String ERROR_CODE_IDENT_REQUIRED = "identRequired";
     public static final String ERROR_CODE_AUTH_REQUIRED = "authRequired";
+    public static final String ERROR_CODE_QUOTA_EXCEEDED = "quotaExceeded";
 
     private List<Object> path;
     private List<SourceLocation> locations;

server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/model/WebServerError.java

@@ -16,7 +16,9 @@
  */
 package io.cloudbeaver.model;
 
+import io.cloudbeaver.DBWebException;
 import org.jkiss.dbeaver.DBException;
+import org.jkiss.dbeaver.model.sql.DBQuotaException;
 
 import java.io.PrintWriter;
 import java.io.StringWriter;
@@ -28,6 +30,7 @@ public class WebServerError {
 
     private String message;
     private String stackTrace;
+    private String errorType;
     private String errorCode;
 
     public WebServerError(Throwable ex) {
@@ -38,6 +41,9 @@ public class WebServerError {
         if (ex instanceof DBException) {
             errorCode = String.valueOf(((DBException) ex).getErrorCode());
         }
+        if (ex instanceof DBQuotaException) {
+            errorType = DBWebException.ERROR_CODE_QUOTA_EXCEEDED;
+        }
     }
 
     public String getMessage() {
@@ -52,4 +58,12 @@ public class WebServerError {
         return errorCode;
     }
 
+    public String getErrorType() {
+        return errorType;
+    }
+
+    public WebServerError getCausedBy() {
+        return null;
+    }
+
 }

server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/server/CBAppConfig.java

@@ -58,6 +58,8 @@ public class CBAppConfig {
     private final Map<String, Object> plugins;
     private final Map<String, AuthProviderConfig> authConfiguration;
 
+    private final Map<String, Object> resourceQuotas;
+
     public CBAppConfig() {
         this.anonymousAccessEnabled = true;
         this.anonymousUserRole = CBConstants.DEFAUL_APP_ANONYMOUS_ROLE_NAME;
@@ -76,6 +78,7 @@ public class CBAppConfig {
         this.defaultNavigatorSettings = DEFAULT_VIEW_SETTINGS;
         this.plugins = new LinkedHashMap<>();
         this.authConfiguration = new LinkedHashMap<>();
+        this.resourceQuotas = new LinkedHashMap<>();
     }
 
     public CBAppConfig(CBAppConfig src) {
@@ -96,6 +99,7 @@ public class CBAppConfig {
         this.defaultNavigatorSettings = src.defaultNavigatorSettings;
         this.plugins = new LinkedHashMap<>(src.plugins);
         this.authConfiguration = new LinkedHashMap<>(src.authConfiguration);
+        this.resourceQuotas = new LinkedHashMap<>(src.resourceQuotas);
     }
 
     public boolean isAnonymousAccessEnabled() {
@@ -236,6 +240,18 @@ public class CBAppConfig {
         return (T)getPluginConfig(pluginId).get(option);
     }
 
+    ////////////////////////////////////////////
+    // Quotas
+
+    public Map<String, Object> getResourceQuotas() {
+        return resourceQuotas;
+    }
+
+    public <T> T getResourceQuota(String quotaId) {
+        return (T) resourceQuotas.get(quotaId);
+    }
+
+
     ////////////////////////////////////////////
     // Auth provider configs
 

server/bundles/io.cloudbeaver.service.data.transfer/src/io/cloudbeaver/service/data/transfer/impl/WebServiceDataTransfer.java

@@ -20,6 +20,7 @@ import io.cloudbeaver.DBWebException;
 import io.cloudbeaver.model.WebAsyncTaskInfo;
 import io.cloudbeaver.model.session.WebAsyncTaskProcessor;
 import io.cloudbeaver.model.session.WebSession;
+import io.cloudbeaver.server.CBApplication;
 import io.cloudbeaver.server.CBPlatform;
 import io.cloudbeaver.service.data.transfer.DBWServiceDataTransfer;
 import io.cloudbeaver.service.sql.WebSQLContextInfo;
@@ -28,9 +29,13 @@ import io.cloudbeaver.service.sql.WebSQLResultsInfo;
 import org.jkiss.code.NotNull;
 import org.jkiss.dbeaver.DBException;
 import org.jkiss.dbeaver.Log;
+import org.jkiss.dbeaver.model.exec.DBCException;
+import org.jkiss.dbeaver.model.exec.DBCResultSet;
+import org.jkiss.dbeaver.model.exec.DBCSession;
 import org.jkiss.dbeaver.model.preferences.DBPPropertyDescriptor;
 import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor;
 import org.jkiss.dbeaver.model.runtime.VoidProgressMonitor;
+import org.jkiss.dbeaver.model.sql.DBQuotaException;
 import org.jkiss.dbeaver.model.struct.DBSDataContainer;
 import org.jkiss.dbeaver.model.struct.DBSEntity;
 import org.jkiss.dbeaver.tools.transfer.IDataTransferConsumer;
@@ -56,6 +61,8 @@ import java.util.stream.Collectors;
  */
 public class WebServiceDataTransfer implements DBWServiceDataTransfer {
 
+    public static final String QUOTE_PROP_FILE_LIMIT = "dataExportFileSizeLimit";
+
     private static final Log log = Log.getLog(WebServiceDataTransfer.class);
 
     private final File dataExportFolder;
@@ -150,6 +157,9 @@ public class WebServiceDataTransfer implements DBWServiceDataTransfer {
                                 log.error("Error deleting export file " + exportFile.getAbsolutePath());
                             }
                         }
+                        if (e instanceof DBException) {
+                            throw e;
+                        }
                         throw new DBException("Error exporting data", e);
                     }
                     WebDataTransferTaskConfig taskConfig = new WebDataTransferTaskConfig(exportFile, parameters);
@@ -181,7 +191,19 @@ public class WebServiceDataTransfer implements DBWServiceDataTransfer {
         }
         IStreamDataExporter exporter = (IStreamDataExporter) processorInstance;
 
-        StreamTransferConsumer consumer = new StreamTransferConsumer();
+        Number fileSizeLimit = CBApplication.getInstance().getAppConfiguration().getResourceQuota(QUOTE_PROP_FILE_LIMIT);
+
+        StreamTransferConsumer consumer = new StreamTransferConsumer() {
+            @Override
+            public void fetchRow(DBCSession session, DBCResultSet resultSet) throws DBCException {
+                super.fetchRow(session, resultSet);
+                if (fileSizeLimit != null && getBytesWritten() > fileSizeLimit.longValue()) {
+                    throw new DBQuotaException(
+                        "Data export quota exceeded", QUOTE_PROP_FILE_LIMIT, fileSizeLimit.longValue(), getBytesWritten());
+                }
+            }
+        };
+
         StreamConsumerSettings settings = new StreamConsumerSettings();
 
         settings.setOutputEncodingBOM(false);
@@ -218,4 +240,5 @@ public class WebServiceDataTransfer implements DBWServiceDataTransfer {
 
         consumer.finishTransfer(monitor, false);
     }
+
 }