exclude columns

Currently only extends things to handle exclude field from PeerDB-io/sqlparser-rs#21 but doesn't do anything with it
PeerDB-io · Nov 3, 2023 · aa72c90 · aa72c90
1 parent 0aec792
commit aa72c90
Show file tree

Hide file tree

Showing 14 changed files with 631 additions and 578 deletions.
diff --git a/flow/activities/flowable.go b/flow/activities/flowable.go
@@ -520,7 +520,6 @@ func (a *FlowableActivity) replicateQRepPartition(ctx context.Context,
 	var stream *model.QRecordStream
 	bufferSize := shared.FetchAndChannelSize
 	var wg sync.WaitGroup
-	var numRecords int64
 
 	var goroutineErr error = nil
 	if config.SourcePeer.Type == protos.DBType_POSTGRES {
@@ -530,7 +529,7 @@ func (a *FlowableActivity) replicateQRepPartition(ctx context.Context,
 		pullPgRecords := func() {
 			pgConn := srcConn.(*connpostgres.PostgresConnector)
 			tmp, err := pgConn.PullQRepRecordStream(config, partition, stream)
-			numRecords = int64(tmp)
+			numRecords := int64(tmp)
 			if err != nil {
 				log.WithFields(log.Fields{
 					"flowName": config.FlowJobName,
@@ -551,7 +550,7 @@ func (a *FlowableActivity) replicateQRepPartition(ctx context.Context,
 		if err != nil {
 			return fmt.Errorf("failed to pull records: %w", err)
 		}
-		numRecords = int64(recordBatch.NumRecords)
+		numRecords := int64(recordBatch.NumRecords)
 		log.WithFields(log.Fields{
 			"flowName": config.FlowJobName,
 		}).Infof("pulled %d records\n", len(recordBatch.Records))

diff --git a/flow/generated/protos/flow.pb.go b/flow/generated/protos/flow.pb.go
diff --git a/flow/model/model.go b/flow/model/model.go
@@ -54,7 +54,7 @@ type RecordItems struct {
 func NewRecordItems() *RecordItems {
 	return &RecordItems{
 		colToValIdx: make(map[string]int),
-		// create a slice of 64 qvalues so that we don't have to allocate memory
+		// create a slice of 32 qvalues so that we don't have to allocate memory
 		// for each record to reduce GC pressure
 		values: make([]*qvalue.QValue, 0, 32),
 	}

diff --git a/nexus/analyzer/src/lib.rs b/nexus/analyzer/src/lib.rs
@@ -164,6 +164,10 @@ impl<'a> StatementAnalyzer for PeerDDLAnalyzer<'a> {
                                     .partition_key
                                     .as_ref()
                                     .map(|s| s.to_string()),
+                                exclude: table_mapping
+                                    .exclude
+                                    .as_ref()
+                                    .map(|ss| ss.iter().map(|s| s.to_string()).collect())
                             });
                         }
 

diff --git a/nexus/flow-rs/src/grpc.rs b/nexus/flow-rs/src/grpc.rs
@@ -170,6 +170,7 @@ impl FlowGrpcClient {
                 source_table_identifier: mapping.source_table_identifier.clone(),
                 destination_table_identifier: mapping.destination_table_identifier.clone(),
                 partition_key: mapping.partition_key.clone().unwrap_or_default(),
+                exclude: mapping.exclude.clone().unwrap_or_default(),
             });
         });
 

diff --git a/nexus/pt/src/flow_model.rs b/nexus/pt/src/flow_model.rs
@@ -10,6 +10,7 @@ pub struct FlowJobTableMapping {
     pub source_table_identifier: String,
     pub destination_table_identifier: String,
     pub partition_key: Option<String>,
+    pub exclude: Option<Vec<String>>,
 }
 
 #[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone)]

diff --git a/nexus/pt/src/peerdb_flow.rs b/nexus/pt/src/peerdb_flow.rs
@@ -36,6 +36,8 @@ pub struct TableMapping {
     pub destination_table_identifier: ::prost::alloc::string::String,
     #[prost(string, tag = "3")]
     pub partition_key: ::prost::alloc::string::String,
+    #[prost(string, repeated, tag = "4")]
+    pub exclude: ::prost::alloc::vec::Vec<::prost::alloc::string::String>,
 }
 #[allow(clippy::derive_partial_eq_without_eq)]
 #[derive(Clone, PartialEq, ::prost::Message)]

diff --git a/nexus/pt/src/peerdb_flow.serde.rs b/nexus/pt/src/peerdb_flow.serde.rs
@@ -5528,6 +5528,9 @@ impl serde::Serialize for TableMapping {
         if !self.partition_key.is_empty() {
             len += 1;
         }
+        if !self.exclude.is_empty() {
+            len += 1;
+        }
         let mut struct_ser = serializer.serialize_struct("peerdb_flow.TableMapping", len)?;
         if !self.source_table_identifier.is_empty() {
             struct_ser.serialize_field("sourceTableIdentifier", &self.source_table_identifier)?;
@@ -5538,6 +5541,9 @@ impl serde::Serialize for TableMapping {
         if !self.partition_key.is_empty() {
             struct_ser.serialize_field("partitionKey", &self.partition_key)?;
         }
+        if !self.exclude.is_empty() {
+            struct_ser.serialize_field("exclude", &self.exclude)?;
+        }
         struct_ser.end()
     }
 }
@@ -5554,13 +5560,15 @@ impl<'de> serde::Deserialize<'de> for TableMapping {
             "destinationTableIdentifier",
             "partition_key",
             "partitionKey",
+            "exclude",
         ];
 
         #[allow(clippy::enum_variant_names)]
         enum GeneratedField {
             SourceTableIdentifier,
             DestinationTableIdentifier,
             PartitionKey,
+            Exclude,
             __SkipField__,
         }
         impl<'de> serde::Deserialize<'de> for GeneratedField {
@@ -5586,6 +5594,7 @@ impl<'de> serde::Deserialize<'de> for TableMapping {
                             "sourceTableIdentifier" | "source_table_identifier" => Ok(GeneratedField::SourceTableIdentifier),
                             "destinationTableIdentifier" | "destination_table_identifier" => Ok(GeneratedField::DestinationTableIdentifier),
                             "partitionKey" | "partition_key" => Ok(GeneratedField::PartitionKey),
+                            "exclude" => Ok(GeneratedField::Exclude),
                             _ => Ok(GeneratedField::__SkipField__),
                         }
                     }
@@ -5608,6 +5617,7 @@ impl<'de> serde::Deserialize<'de> for TableMapping {
                 let mut source_table_identifier__ = None;
                 let mut destination_table_identifier__ = None;
                 let mut partition_key__ = None;
+                let mut exclude__ = None;
                 while let Some(k) = map.next_key()? {
                     match k {
                         GeneratedField::SourceTableIdentifier => {
@@ -5628,6 +5638,12 @@ impl<'de> serde::Deserialize<'de> for TableMapping {
                             }
                             partition_key__ = Some(map.next_value()?);
                         }
+                        GeneratedField::Exclude => {
+                            if exclude__.is_some() {
+                                return Err(serde::de::Error::duplicate_field("exclude"));
+                            }
+                            exclude__ = Some(map.next_value()?);
+                        }
                         GeneratedField::__SkipField__ => {
                             let _ = map.next_value::<serde::de::IgnoredAny>()?;
                         }
@@ -5637,6 +5653,7 @@ impl<'de> serde::Deserialize<'de> for TableMapping {
                     source_table_identifier: source_table_identifier__.unwrap_or_default(),
                     destination_table_identifier: destination_table_identifier__.unwrap_or_default(),
                     partition_key: partition_key__.unwrap_or_default(),
+                    exclude: exclude__.unwrap_or_default(),
                 })
             }
         }

diff --git a/nexus/sqlparser-rs b/nexus/sqlparser-rs
diff --git a/protos/flow.proto b/protos/flow.proto
@@ -26,6 +26,7 @@ message TableMapping {
   string source_table_identifier = 1;
   string destination_table_identifier = 2;
   string partition_key = 3;
+  repeated string exclude = 4;
 }
 
 message FlowConnectionConfigs {

diff --git a/ui/app/dto/MirrorsDTO.ts b/ui/app/dto/MirrorsDTO.ts
@@ -17,4 +17,5 @@ export type TableMapRow = {
   source: string;
   destination: string;
   partitionKey: string;
+  exclude: string[];
 };
diff --git a/ui/app/mirrors/create/handlers.ts b/ui/app/mirrors/create/handlers.ts
@@ -58,6 +58,7 @@ const reformattedTableMapping = (tableMapping: TableMapRow[]) => {
       sourceTableIdentifier: row.source,
       destinationTableIdentifier: row.destination,
       partitionKey: row.partitionKey,
+      exclude: row.exclude,
     };
   });
   return mapping;

diff --git a/ui/app/mirrors/create/tablemapping.tsx b/ui/app/mirrors/create/tablemapping.tsx
@@ -32,7 +32,7 @@ const TableMapping = ({
   const [loading, setLoading] = useState(false);
 
   const handleAddRow = (source: string) => {
-    setRows([...rows, { source, destination: source, partitionKey: '' }]);
+    setRows([...rows, { source, destination: source, partitionKey: '', exclude: [] }]);
   };
 
   const handleRemoveRow = (source: string) => {

diff --git a/ui/grpc_generated/flow.ts b/ui/grpc_generated/flow.ts
@@ -101,6 +101,7 @@ export interface TableMapping {
   sourceTableIdentifier: string;
   destinationTableIdentifier: string;
   partitionKey: string;
+  exclude: string[];
 }
 
 export interface FlowConnectionConfigs {
@@ -691,7 +692,7 @@ export const RelationMessage = {
 };
 
 function createBaseTableMapping(): TableMapping {
-  return { sourceTableIdentifier: "", destinationTableIdentifier: "", partitionKey: "" };
+  return { sourceTableIdentifier: "", destinationTableIdentifier: "", partitionKey: "", exclude: [] };
 }
 
 export const TableMapping = {
@@ -705,6 +706,9 @@ export const TableMapping = {
     if (message.partitionKey !== "") {
       writer.uint32(26).string(message.partitionKey);
     }
+    for (const v of message.exclude) {
+      writer.uint32(34).string(v!);
+    }
     return writer;
   },
 
@@ -736,6 +740,13 @@ export const TableMapping = {
 
           message.partitionKey = reader.string();
           continue;
+        case 4:
+          if (tag !== 34) {
+            break;
+          }
+
+          message.exclude.push(reader.string());
+          continue;
       }
       if ((tag & 7) === 4 || tag === 0) {
         break;
@@ -752,6 +763,7 @@ export const TableMapping = {
         ? String(object.destinationTableIdentifier)
         : "",
       partitionKey: isSet(object.partitionKey) ? String(object.partitionKey) : "",
+      exclude: Array.isArray(object?.exclude) ? object.exclude.map((e: any) => String(e)) : [],
     };
   },
 
@@ -766,6 +778,9 @@ export const TableMapping = {
     if (message.partitionKey !== "") {
       obj.partitionKey = message.partitionKey;
     }
+    if (message.exclude?.length) {
+      obj.exclude = message.exclude;
+    }
     return obj;
   },
 
@@ -777,6 +792,7 @@ export const TableMapping = {
     message.sourceTableIdentifier = object.sourceTableIdentifier ?? "";
     message.destinationTableIdentifier = object.destinationTableIdentifier ?? "";
     message.partitionKey = object.partitionKey ?? "";
+    message.exclude = object.exclude?.map((e) => e) || [];
     return message;
   },
 };
+13 −0		src/ast/mod.rs
+30 −1		src/parser.rs
+12 −2		tests/sqlparser_postgres.rs