s3: avro codec field

Author: serprex

flow/connectors/s3/qrep.go

@@ -68,7 +68,21 @@ func (c *S3Connector) writeToAvroFile(
 
 	s3AvroFileKey := fmt.Sprintf("%s/%s/%s.avro", s3o.Prefix, jobName, partitionID)
 
-	writer := avro.NewPeerDBOCFWriter(stream, avroSchema, ocf.Null, protos.DBType_S3)
+	var codec ocf.CodecName
+	switch c.codec {
+	case protos.AvroCodec_Null:
+		codec = ocf.Null
+	case protos.AvroCodec_Deflate:
+		codec = ocf.Deflate
+	case protos.AvroCodec_Snappy:
+		codec = ocf.Snappy
+	case protos.AvroCodec_ZStandard:
+		codec = ocf.ZStandard
+	default:
+		return 0, fmt.Errorf("unsupported codec %s", c.codec)
+	}
+
+	writer := avro.NewPeerDBOCFWriter(stream, avroSchema, codec, protos.DBType_S3)
 	avroFile, err := writer.WriteRecordsToS3(ctx, env, s3o.Bucket, s3AvroFileKey, c.credentialsProvider, nil, nil)
 	if err != nil {
 		return 0, fmt.Errorf("failed to write records to S3: %w", err)

flow/connectors/s3/s3.go

@@ -21,6 +21,7 @@ type S3Connector struct {
 	credentialsProvider utils.AWSCredentialsProvider
 	client              s3.Client
 	url                 string
+	codec               protos.AvroCodec
 }
 
 func NewS3Connector(
@@ -44,11 +45,12 @@ func NewS3Connector(
 		return nil, err
 	}
 	return &S3Connector{
-		url:                 config.Url,
 		PostgresMetadata:    pgMetadata,
 		client:              *s3Client,
 		credentialsProvider: provider,
 		logger:              logger,
+		url:                 config.Url,
+		codec:               config.Codec,
 	}, nil
 }
 

nexus/analyzer/src/lib.rs

@@ -731,6 +731,11 @@ fn parse_db_options(db_type: DbType, with_options: &[SqlOption]) -> anyhow::Resu
                     .get("tls_host")
                     .map(|s| s.to_string())
                     .unwrap_or_default(),
+                codec: opts
+                    .get("codec")
+                    .and_then(|s| pt::peerdb_peers::AvroCodec::from_str_name(s))
+                    .map(|codec| codec.into())
+                    .unwrap_or_default(),
             };
             Config::S3Config(s3_config)
         }

protos/peers.proto

@@ -134,6 +134,13 @@ message EventHubGroupConfig {
   repeated string unnest_columns = 3;
 }
 
+enum AvroCodec {
+  Null = 0;
+  Deflate = 1;
+  Snappy = 2;
+  ZStandard = 3;
+}
+
 message S3Config {
   string url = 1;
   optional string access_key_id = 2 [(peerdb_redacted) = true];
@@ -143,6 +150,7 @@ message S3Config {
   optional string endpoint = 6;
   optional string root_ca = 7 [(peerdb_redacted) = true];
   string tls_host = 8;
+  AvroCodec codec = 9;
 }
 
 message ClickhouseConfig{

ui/app/peers/create/[peerType]/helpers/s3.ts

@@ -1,4 +1,4 @@
-import { S3Config } from '@/grpc_generated/peers';
+import { AvroCodec, S3Config, avroCodecFromJSON } from '@/grpc_generated/peers';
 import { PeerSetting } from './common';
 
 export const s3Setting: PeerSetting[] = [
@@ -74,6 +74,20 @@ export const s3Setting: PeerSetting[] = [
     tips: 'Overrides expected hostname during tls cert verification.',
     optional: true,
   },
+  {
+    label: 'Avro Codec',
+    field: 'codec',
+    stateHandler: (value, setter) =>
+      setter((curr) => ({ ...curr, codec: avroCodecFromJSON(value) })),
+    type: 'select',
+    placeholder: 'Select avro codec',
+    options: [
+      { value: 'Null', label: 'Null' },
+      { value: 'Deflate', label: 'Deflate' },
+      { value: 'Snappy', label: 'Snappy' },
+      { value: 'ZStandard', label: 'ZStandard' },
+    ],
+  },
 ];
 
 export const blankS3Setting: S3Config = {
@@ -85,4 +99,5 @@ export const blankS3Setting: S3Config = {
   endpoint: '',
   rootCa: undefined,
   tlsHost: '',
+  codec: AvroCodec.Null,
 };

ui/app/peers/create/[peerType]/schema.ts

@@ -1,5 +1,6 @@
 import { ehSchema } from '@/components/PeerForms/Eventhubs/schema';
 import {
+  AvroCodec,
   ElasticsearchAuthType,
   MySqlFlavor,
   MySqlReplicationMechanism,
@@ -526,6 +527,12 @@ export const s3Schema = z.object({
       error: () => 'Endpoint must be a string',
     })
     .optional(),
+  codec: z.enum(AvroCodec, {
+    error: (issue) =>
+      issue.input === undefined
+        ? 'Avro codec is required'
+        : 'Avro codec must be one of [Null,Deflate,Snappy,ZStandard]',
+  }),
 });
 
 export const psSchema = z.object({

ui/components/PeerForms/S3Form.tsx

@@ -1,14 +1,20 @@
 'use client';
 import { PeerSetter } from '@/app/dto/PeersDTO';
 import { s3Setting } from '@/app/peers/create/[peerType]/helpers/s3';
+import SelectTheme from '@/app/styles/select';
 import { GCS_ENDPOINT } from '@/app/utils/gcsEndpoint';
 import InfoPopover from '@/components/InfoPopover';
 import { Label } from '@/lib/Label';
-import { RowWithRadiobutton, RowWithTextField } from '@/lib/Layout';
+import {
+  RowWithRadiobutton,
+  RowWithSelect,
+  RowWithTextField,
+} from '@/lib/Layout';
 import { RadioButton, RadioButtonGroup } from '@/lib/RadioButtonGroup';
 import { TextField } from '@/lib/TextField';
 import { Tooltip } from '@/lib/Tooltip';
 import { useEffect, useState } from 'react';
+import ReactSelect from 'react-select';
 import { handleFieldChange } from './common';
 
 interface S3Props {
@@ -65,9 +71,25 @@ export default function S3Form({ setter }: S3Props) {
           action={<RadioButton value='GCS' />}
         />
       </RadioButtonGroup>
-      {s3Setting.map((setting, index) => {
-        if (displayCondition(setting.label))
-          return (
+      {s3Setting.map(
+        (setting, index) =>
+          displayCondition(setting.label) &&
+          (setting.type === 'select' ? (
+            <RowWithSelect
+              key={index}
+              label={<Label>{setting.label}</Label>}
+              action={
+                <ReactSelect
+                  placeholder={setting.placeholder}
+                  onChange={(val) =>
+                    val && setting.stateHandler(val.value, setter)
+                  }
+                  options={setting.options}
+                  theme={SelectTheme}
+                />
+              }
+            />
+          ) : (
             <RowWithTextField
               key={index}
               label={
@@ -115,8 +137,8 @@ export default function S3Form({ setter }: S3Props) {
                 </div>
               }
             />
-          );
-      })}
+          ))
+      )}
     </div>
   );
 }