[parquet-sdk]migrate objects [pt2]

rust/processor/src/db/common/models/object_models/raw_v2_objects.rs

@@ -9,7 +9,10 @@ use super::v2_object_utils::{CurrentObjectPK, ObjectAggregatedDataMapping};
 use crate::{
     db::postgres::models::default_models::move_resources::MoveResource,
     schema::current_objects,
-    utils::{database::DbPoolConnection, util::standardize_address},
+    utils::{
+        database::{DbContext, DbPoolConnection},
+        util::standardize_address,
+    },
 };
 use ahash::AHashMap;
 use aptos_protos::transaction::v1::{DeleteResource, WriteResource};
@@ -19,6 +22,8 @@ use diesel_async::RunQueryDsl;
 use field_count::FieldCount;
 use serde::{Deserialize, Serialize};
 
+const DELETED_RESOURCE_OWNER_ADDRESS: &str = "Unknown";
+
 #[derive(Clone, Debug, Deserialize, FieldCount, Serialize)]
 pub struct RawObject {
     pub transaction_version: i64,
@@ -30,6 +35,7 @@ pub struct RawObject {
     pub allow_ungated_transfer: bool,
     pub is_deleted: bool,
     pub untransferrable: bool,
+    pub block_timestamp: chrono::NaiveDateTime,
 }
 
 pub trait ObjectConvertible {
@@ -46,6 +52,7 @@ pub struct RawCurrentObject {
     pub last_transaction_version: i64,
     pub is_deleted: bool,
     pub untransferrable: bool,
+    pub block_timestamp: chrono::NaiveDateTime,
 }
 
 pub trait CurrentObjectConvertible {
@@ -73,14 +80,10 @@ impl RawObject {
         txn_version: i64,
         write_set_change_index: i64,
         object_metadata_mapping: &ObjectAggregatedDataMapping,
+        block_timestamp: chrono::NaiveDateTime,
     ) -> anyhow::Result<Option<(Self, RawCurrentObject)>> {
         let address = standardize_address(&write_resource.address.to_string());
-        println!("address: {:?}", address);
         if let Some(object_aggregated_metadata) = object_metadata_mapping.get(&address) {
-            println!(
-                "object_aggregated_metadata: {:?}",
-                object_aggregated_metadata
-            );
             // do something
             let object_with_metadata = object_aggregated_metadata.object.clone();
             let object_core = object_with_metadata.object_core;
@@ -101,6 +104,7 @@ impl RawObject {
                     allow_ungated_transfer: object_core.allow_ungated_transfer,
                     is_deleted: false,
                     untransferrable,
+                    block_timestamp,
                 },
                 RawCurrentObject {
                     object_address: address,
@@ -111,6 +115,7 @@ impl RawObject {
                     last_transaction_version: txn_version,
                     is_deleted: false,
                     untransferrable,
+                    block_timestamp,
                 },
             )))
         } else {
@@ -126,9 +131,8 @@ impl RawObject {
         txn_version: i64,
         write_set_change_index: i64,
         object_mapping: &AHashMap<CurrentObjectPK, RawCurrentObject>,
-        conn: &mut DbPoolConnection<'_>,
-        query_retries: u32,
-        query_retry_delay_ms: u64,
+        db_context: &mut Option<DbContext<'_>>,
+        block_timestamp: chrono::NaiveDateTime,
     ) -> anyhow::Result<Option<(Self, RawCurrentObject)>> {
         if delete_resource.type_str == "0x1::object::ObjectGroup" {
             let resource = MoveResource::from_delete_resource(
@@ -137,52 +141,91 @@ impl RawObject {
                 txn_version,
                 0, // Placeholder, this isn't used anyway
             );
-            let previous_object = if let Some(object) = object_mapping.get(&resource.address) {
-                object.clone()
+            // Add a logid here to handle None conn
+            if db_context.is_none() {
+                // This is a hack to prevent the program for parquet
+                Ok(Some((
+                    Self {
+                        transaction_version: txn_version,
+                        write_set_change_index,
+                        object_address: resource.address.clone(),
+                        owner_address: DELETED_RESOURCE_OWNER_ADDRESS.to_string(),
+                        state_key_hash: resource.state_key_hash.clone(),
+                        guid_creation_num: BigDecimal::default(),
+                        allow_ungated_transfer: false,
+                        is_deleted: true,
+                        untransferrable: false,
+                        block_timestamp,
+                    },
+                    RawCurrentObject {
+                        object_address: resource.address.clone(),
+                        owner_address: DELETED_RESOURCE_OWNER_ADDRESS.to_string(),
+                        state_key_hash: resource.state_key_hash.clone(),
+                        last_guid_creation_num: BigDecimal::default(),
+                        allow_ungated_transfer: false,
+                        last_transaction_version: txn_version,
+                        is_deleted: true,
+                        untransferrable: false,
+                        block_timestamp,
+                    },
+                )))
             } else {
-                match Self::get_current_object(
-                    conn,
-                    &resource.address,
-                    query_retries,
-                    query_retry_delay_ms,
-                )
-                .await
-                {
-                    Ok(object) => object,
-                    Err(_) => {
-                        tracing::error!(
+                let previous_object = if let Some(object) = object_mapping.get(&resource.address) {
+                    object.clone()
+                } else if let Some(db_context) = db_context {
+                    match Self::get_current_object(
+                        &mut db_context.conn,
+                        &resource.address,
+                        db_context.query_retries,
+                        db_context.query_retry_delay_ms,
+                    )
+                    .await
+                    {
+                        Ok(object) => object,
+                        Err(_) => {
+                            tracing::error!(
                             transaction_version = txn_version,
                             lookup_key = &resource.address,
                             "Missing current_object for object_address: {}. You probably should backfill db.",
                             resource.address,
                         );
-                        return Ok(None);
+                            return Ok(None);
+                        },
+                    }
+                } else {
+                    tracing::error!(
+                        transaction_version = txn_version,
+                        lookup_key = &resource.address,
+                        "Connection to DB is missing. You may need to investigate",
+                    );
+                    return Ok(None);
+                };
+                Ok(Some((
+                    Self {
+                        transaction_version: txn_version,
+                        write_set_change_index,
+                        object_address: resource.address.clone(),
+                        owner_address: previous_object.owner_address.clone(),
+                        state_key_hash: resource.state_key_hash.clone(),
+                        guid_creation_num: previous_object.last_guid_creation_num.clone(),
+                        allow_ungated_transfer: previous_object.allow_ungated_transfer,
+                        is_deleted: true,
+                        untransferrable: previous_object.untransferrable,
+                        block_timestamp,
                     },
-                }
-            };
-            Ok(Some((
-                Self {
-                    transaction_version: txn_version,
-                    write_set_change_index,
-                    object_address: resource.address.clone(),
-                    owner_address: previous_object.owner_address.clone(),
-                    state_key_hash: resource.state_key_hash.clone(),
-                    guid_creation_num: previous_object.last_guid_creation_num.clone(),
-                    allow_ungated_transfer: previous_object.allow_ungated_transfer,
-                    is_deleted: true,
-                    untransferrable: previous_object.untransferrable,
-                },
-                RawCurrentObject {
-                    object_address: resource.address,
-                    owner_address: previous_object.owner_address.clone(),
-                    state_key_hash: resource.state_key_hash,
-                    last_guid_creation_num: previous_object.last_guid_creation_num.clone(),
-                    allow_ungated_transfer: previous_object.allow_ungated_transfer,
-                    last_transaction_version: txn_version,
-                    is_deleted: true,
-                    untransferrable: previous_object.untransferrable,
-                },
-            )))
+                    RawCurrentObject {
+                        object_address: resource.address,
+                        owner_address: previous_object.owner_address.clone(),
+                        state_key_hash: resource.state_key_hash,
+                        last_guid_creation_num: previous_object.last_guid_creation_num.clone(),
+                        allow_ungated_transfer: previous_object.allow_ungated_transfer,
+                        last_transaction_version: txn_version,
+                        is_deleted: true,
+                        untransferrable: previous_object.untransferrable,
+                        block_timestamp,
+                    },
+                )))
+            }
         } else {
             Ok(None)
         }
@@ -210,6 +253,7 @@ impl RawObject {
                         last_transaction_version: res.last_transaction_version,
                         is_deleted: res.is_deleted,
                         untransferrable: res.untransferrable,
+                        block_timestamp: chrono::NaiveDateTime::default(), // this won't be used
                     });
                 },
                 Err(_) => {

rust/processor/src/db/common/models/token_v2_models/raw_v2_token_ownerships.rs

@@ -21,7 +21,7 @@ use crate::{
     },
     schema::current_token_ownerships_v2,
     utils::{
-        database::DbPoolConnection,
+        database::{DbContext, DbPoolConnection},
         util::{ensure_not_negative, standardize_address},
     },
 };
@@ -259,9 +259,7 @@ impl RawTokenOwnershipV2 {
         prior_nft_ownership: &AHashMap<String, NFTOwnershipV2>,
         tokens_burned: &TokenV2Burned,
         object_metadatas: &ObjectAggregatedDataMapping,
-        conn: &mut Option<DbPoolConnection<'_>>,
-        query_retries: u32,
-        query_retry_delay_ms: u64,
+        db_context: &mut Option<DbContext<'_>>,
     ) -> anyhow::Result<Option<(Self, RawCurrentTokenOwnershipV2)>> {
         let token_data_id = standardize_address(&write_resource.address.to_string());
         if tokens_burned
@@ -328,9 +326,7 @@ impl RawTokenOwnershipV2 {
                     txn_timestamp,
                     prior_nft_ownership,
                     tokens_burned,
-                    conn,
-                    query_retries,
-                    query_retry_delay_ms,
+                    db_context,
                 )
                 .await;
             }
@@ -346,9 +342,7 @@ impl RawTokenOwnershipV2 {
         txn_timestamp: chrono::NaiveDateTime,
         prior_nft_ownership: &AHashMap<String, NFTOwnershipV2>,
         tokens_burned: &TokenV2Burned,
-        conn: &mut Option<DbPoolConnection<'_>>,
-        query_retries: u32,
-        query_retry_delay_ms: u64,
+        db_context: &mut Option<DbContext<'_>>,
     ) -> anyhow::Result<Option<(Self, RawCurrentTokenOwnershipV2)>> {
         let token_address = standardize_address(&delete_resource.address.to_string());
         Self::get_burned_nft_v2_helper(
@@ -358,9 +352,7 @@ impl RawTokenOwnershipV2 {
             txn_timestamp,
             prior_nft_ownership,
             tokens_burned,
-            conn,
-            query_retries,
-            query_retry_delay_ms,
+            db_context,
         )
         .await
     }
@@ -372,9 +364,7 @@ impl RawTokenOwnershipV2 {
         txn_timestamp: chrono::NaiveDateTime,
         prior_nft_ownership: &AHashMap<String, NFTOwnershipV2>,
         tokens_burned: &TokenV2Burned,
-        conn: &mut Option<DbPoolConnection<'_>>,
-        query_retries: u32,
-        query_retry_delay_ms: u64,
+        db_context: &mut Option<DbContext<'_>>,
     ) -> anyhow::Result<Option<(Self, RawCurrentTokenOwnershipV2)>> {
         let token_address = standardize_address(token_address);
         if let Some(burn_event) = tokens_burned.get(&token_address) {
@@ -389,7 +379,7 @@ impl RawTokenOwnershipV2 {
                 match prior_nft_ownership.get(&token_address) {
                     Some(inner) => inner.owner_address.clone(),
                     None => {
-                        match conn {
+                        match db_context {
                             None => {
                                 // TODO: update message
                                 tracing::error!(
@@ -399,12 +389,12 @@ impl RawTokenOwnershipV2 {
                                 );
                                 DEFAULT_OWNER_ADDRESS.to_string()
                             },
-                            Some(ref mut conn) => {
+                            Some(db_context) => {
                                 match CurrentTokenOwnershipV2Query::get_latest_owned_nft_by_token_data_id(
-                                    conn,
+                                    &mut db_context.conn,
                                     &token_address,
-                                    query_retries,
-                                    query_retry_delay_ms,
+                                    db_context.query_retries,
+                                    db_context.query_retry_delay_ms,
                                 )
                                     .await
                                 {

rust/processor/src/db/parquet/models/mod.rs

@@ -3,6 +3,7 @@ pub mod ans_models;
 pub mod default_models;
 pub mod event_models;
 pub mod fungible_asset_models;
+pub mod object_models;
 pub mod token_v2_models;
 pub mod transaction_metadata_model;
 pub mod user_transaction_models;