Merge pull request #25123 from petrosagg/derive-arbitrary

storage: derive `Arbitrary` where possible

Author: petrosagg

Cargo.lock

@@ -22,6 +22,7 @@ mysql_async = { version = "0.33.0", default-features = false, features = ["minim
 once_cell = "1.16.0"
 prost = { version = "0.11.3", features = ["no-recursion-limit"] }
 proptest = { version = "1.0.0", default-features = false, features = ["std"] }
+proptest-derive = { version = "0.3.0", features = ["boxed_union"]}
 serde = { version = "1.0.152", features = ["derive"] }
 thiserror = "1.0.37"
 tracing = "0.1.37"

src/mysql-util/Cargo.toml

@@ -11,13 +11,13 @@ use std::collections::BTreeSet;
 
 use mz_proto::{IntoRustIfSome, RustType, TryFromProtoError};
 use mz_repr::ColumnType;
-use proptest::prelude::{any, Arbitrary};
-use proptest::strategy::{BoxedStrategy, Strategy};
+use proptest::prelude::any;
+use proptest_derive::Arbitrary;
 use serde::{Deserialize, Serialize};
 
 include!(concat!(env!("OUT_DIR"), "/mz_mysql_util.rs"));
 
-#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Arbitrary)]
 pub struct MySqlTableDesc {
     /// In MySQL the schema and database of a table are synonymous.
     pub schema_name: String,
@@ -28,9 +28,11 @@ pub struct MySqlTableDesc {
     /// The index of each column is based on its `ordinal_position`
     /// reported by the information_schema.columns table, which defines
     /// the order of column values when received in a row.
+    #[proptest(strategy = "proptest::collection::vec(any::<MySqlColumnDesc>(), 0..4)")]
     pub columns: Vec<MySqlColumnDesc>,
     /// Applicable keys for this table (i.e. primary key and unique
     /// constraints).
+    #[proptest(strategy = "proptest::collection::btree_set(any::<MySqlKeyDesc>(), 0..4)")]
     pub keys: BTreeSet<MySqlKeyDesc>,
 }
 
@@ -62,28 +64,7 @@ impl RustType<ProtoMySqlTableDesc> for MySqlTableDesc {
     }
 }
 
-impl Arbitrary for MySqlTableDesc {
-    type Parameters = ();
-    type Strategy = BoxedStrategy<Self>;
-
-    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
-        (
-            any::<String>(),
-            any::<String>(),
-            proptest::collection::vec(any::<MySqlColumnDesc>(), 1..4),
-            proptest::collection::btree_set(any::<MySqlKeyDesc>(), 1..4),
-        )
-            .prop_map(|(schema_name, name, columns, keys)| Self {
-                schema_name,
-                name,
-                columns,
-                keys,
-            })
-            .boxed()
-    }
-}
-
-#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Arbitrary)]
 pub struct MySqlColumnDesc {
     /// The name of the column.
     pub name: String,
@@ -109,24 +90,14 @@ impl RustType<ProtoMySqlColumnDesc> for MySqlColumnDesc {
     }
 }
 
-impl Arbitrary for MySqlColumnDesc {
-    type Parameters = ();
-    type Strategy = BoxedStrategy<Self>;
-
-    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
-        (any::<String>(), any::<ColumnType>())
-            .prop_map(|(name, column_type)| Self { name, column_type })
-            .boxed()
-    }
-}
-
-#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Ord, PartialOrd)]
+#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Ord, PartialOrd, Arbitrary)]
 pub struct MySqlKeyDesc {
     /// The name of the index.
     pub name: String,
     /// Whether or not this key is the primary key.
     pub is_primary: bool,
     /// The columns that make up the key.
+    #[proptest(strategy = "proptest::collection::vec(any::<String>(), 0..4)")]
     pub columns: Vec<String>,
 }
 
@@ -147,18 +118,3 @@ impl RustType<ProtoMySqlKeyDesc> for MySqlKeyDesc {
         })
     }
 }
-
-impl Arbitrary for MySqlKeyDesc {
-    type Parameters = ();
-    type Strategy = BoxedStrategy<Self>;
-
-    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
-        (any::<String>(), any::<bool>(), any::<Vec<String>>())
-            .prop_map(|(name, is_primary, columns)| Self {
-                name,
-                is_primary,
-                columns,
-            })
-            .boxed()
-    }
-}

src/mysql-util/src/desc.rs

@@ -26,6 +26,7 @@ postgres-openssl = { version = "0.5.0" }
 proptest = { version = "1.0.0", default-features = false, features = [
     "std",
 ], optional = true }
+proptest-derive = { version = "0.3.0", features = ["boxed_union"]}
 prost = { version = "0.11.3", features = [
     "no-recursion-limit",
 ], optional = true }

src/postgres-util/Cargo.toml

@@ -13,8 +13,8 @@ use std::collections::BTreeSet;
 
 use anyhow::bail;
 use mz_proto::{IntoRustIfSome, RustType, TryFromProtoError};
-use proptest::prelude::{any, Arbitrary};
-use proptest::strategy::{BoxedStrategy, Strategy};
+use proptest::prelude::any;
+use proptest_derive::Arbitrary;
 use serde::{Deserialize, Serialize};
 use tokio_postgres::types::Oid;
 use tracing::warn;
@@ -35,7 +35,7 @@ pub struct PostgresSchemaDesc {
 }
 
 /// Describes a table in a PostgreSQL database.
-#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Arbitrary)]
 pub struct PostgresTableDesc {
     /// The OID of the table.
     pub oid: Oid,
@@ -44,9 +44,11 @@ pub struct PostgresTableDesc {
     /// The name of the table.
     pub name: String,
     /// The description of each column, in order of their position in the table.
+    #[proptest(strategy = "proptest::collection::vec(any::<PostgresColumnDesc>(), 1..4)")]
     pub columns: Vec<PostgresColumnDesc>,
     /// Applicable keys for this table (i.e. primary key and unique
     /// constraints).
+    #[proptest(strategy = "proptest::collection::btree_set(any::<PostgresKeyDesc>(), 1..4)")]
     pub keys: BTreeSet<PostgresKeyDesc>,
 }
 
@@ -128,31 +130,8 @@ impl RustType<ProtoPostgresTableDesc> for PostgresTableDesc {
     }
 }
 
-impl Arbitrary for PostgresTableDesc {
-    type Strategy = BoxedStrategy<Self>;
-    type Parameters = ();
-
-    fn arbitrary_with(_: Self::Parameters) -> Self::Strategy {
-        (
-            any::<String>(),
-            any::<String>(),
-            any::<u32>(),
-            proptest::collection::vec(any::<PostgresColumnDesc>(), 1..4),
-            proptest::collection::btree_set(any::<PostgresKeyDesc>(), 1..4),
-        )
-            .prop_map(|(name, namespace, oid, columns, keys)| PostgresTableDesc {
-                name,
-                namespace,
-                oid,
-                columns,
-                keys,
-            })
-            .boxed()
-    }
-}
-
 /// Describes a column in a [`PostgresTableDesc`].
-#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
+#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, Arbitrary)]
 pub struct PostgresColumnDesc {
     /// The name of the column.
     pub name: String,
@@ -214,40 +193,16 @@ impl RustType<ProtoPostgresColumnDesc> for PostgresColumnDesc {
     }
 }
 
-impl Arbitrary for PostgresColumnDesc {
-    type Strategy = BoxedStrategy<Self>;
-    type Parameters = ();
-
-    fn arbitrary_with(_: Self::Parameters) -> Self::Strategy {
-        (
-            any::<String>(),
-            any::<u16>(),
-            any::<u32>(),
-            any::<i32>(),
-            any::<bool>(),
-        )
-            .prop_map(
-                |(name, col_num, type_oid, type_mod, nullable)| PostgresColumnDesc {
-                    name,
-                    col_num,
-                    type_oid,
-                    type_mod,
-                    nullable,
-                },
-            )
-            .boxed()
-    }
-}
-
 /// Describes a key in a [`PostgresTableDesc`].
-#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, PartialOrd, Ord)]
+#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize, PartialOrd, Ord, Arbitrary)]
 pub struct PostgresKeyDesc {
     /// This key is derived from the `pg_constraint` with this OID.
     pub oid: Oid,
     /// The name of the constraints.
     pub name: String,
     /// The `attnum` of the columns comprising the key. `attnum` is a unique identifier for a column
     /// in a PG table; see <https://www.postgresql.org/docs/current/catalog-pg-attribute.html>
+    #[proptest(strategy = "proptest::collection::vec(any::<u16>(), 0..4)")]
     pub cols: Vec<u16>,
     /// Whether or not this key is the primary key.
     pub is_primary: bool,
@@ -281,28 +236,3 @@ impl RustType<ProtoPostgresKeyDesc> for PostgresKeyDesc {
         })
     }
 }
-
-impl Arbitrary for PostgresKeyDesc {
-    type Strategy = BoxedStrategy<Self>;
-    type Parameters = ();
-
-    fn arbitrary_with(_: Self::Parameters) -> Self::Strategy {
-        (
-            any::<u32>(),
-            any::<String>(),
-            any::<Vec<u16>>(),
-            any::<bool>(),
-            any::<bool>(),
-        )
-            .prop_map(
-                |(oid, name, cols, is_primary, nulls_not_distinct)| PostgresKeyDesc {
-                    oid,
-                    name,
-                    cols,
-                    is_primary,
-                    nulls_not_distinct,
-                },
-            )
-            .boxed()
-    }
-}

src/postgres-util/src/desc.rs

@@ -30,6 +30,7 @@ mz-service = { path = "../service" }
 mz-storage-types = { path = "../storage-types" }
 mz-timely-util = { path = "../timely-util" }
 proptest = { version = "1.0.0", default-features = false, features = ["std"] }
+proptest-derive = { version = "0.3.0", features = ["boxed_union"]}
 prometheus = { version = "0.13.3", default-features = false }
 prost = { version = "0.11.3", features = ["no-recursion-limit"] }
 rdkafka = { version = "0.29.0", features = [