55
66 "gorm.io/gorm"
77 "gorm.io/gorm/clause"
8+ "k8s.io/utils/clock"
89
910 "github.com/flyteorg/flyte/datacatalog/pkg/common"
1011 "github.com/flyteorg/flyte/datacatalog/pkg/repositories/errors"
@@ -18,13 +19,20 @@ type artifactRepo struct {
1819 db * gorm.DB
1920 errorTransformer errors.ErrorTransformer
2021 repoMetrics gormMetrics
22+ clock clock.Clock
2123}
2224
23- func NewArtifactRepo (db * gorm.DB , errorTransformer errors.ErrorTransformer , scope promutils.Scope ) interfaces.ArtifactRepo {
25+ func NewArtifactRepo (
26+ db * gorm.DB ,
27+ errorTransformer errors.ErrorTransformer ,
28+ scope promutils.Scope ,
29+ clock clock.Clock ,
30+ ) interfaces.ArtifactRepo {
2431 return & artifactRepo {
2532 db : db ,
2633 errorTransformer : errorTransformer ,
2734 repoMetrics : newGormMetrics (scope ),
35+ clock : clock ,
2836 }
2937}
3038
@@ -35,6 +43,19 @@ func (h *artifactRepo) Create(ctx context.Context, artifact models.Artifact) err
3543
3644 tx := h .db .WithContext (ctx ).Begin ()
3745
46+ existing , err := h .getForUpdate (ctx , artifact .ArtifactKey , tx )
47+ if err == nil {
48+ if existing .ExpiresAt != nil && h .clock .Now ().After (* existing .ExpiresAt ) {
49+ // If the previous artifact is expired, soft delete it before creating a new record
50+ tx .Delete (& existing )
51+ }
52+ } else {
53+ if err .Error () != gorm .ErrRecordNotFound .Error () {
54+ return h .errorTransformer .ToDataCatalogError (err )
55+ }
56+ // Not found is desirable
57+ }
58+
3859 tx = tx .Create (& artifact )
3960
4061 if tx .Error != nil {
@@ -50,17 +71,19 @@ func (h *artifactRepo) Create(ctx context.Context, artifact models.Artifact) err
5071 return nil
5172}
5273
53- func (h * artifactRepo ) Get (ctx context.Context , in models.ArtifactKey ) (models.Artifact , error ) {
74+ func (h * artifactRepo ) GetAndFilterExpired (ctx context.Context , in models.ArtifactKey ) (models.Artifact , error ) {
5475 timer := h .repoMetrics .GetDuration .Start (ctx )
5576 defer timer .Stop ()
5677
5778 var artifact models.Artifact
58- result := h .db .WithContext (ctx ).Preload ("ArtifactData" ).
79+ result := h .db .WithContext (ctx ).
80+ Where ("artifacts.expires_at is null or artifacts.expires_at < ?" , h .clock .Now ()).
81+ Preload ("ArtifactData" ).
5982 Preload ("Partitions" , func (db * gorm.DB ) * gorm.DB {
6083 return db .WithContext (ctx ).Order ("partitions.created_at ASC" ) // preserve the order in which the partitions were created
6184 }).
6285 Preload ("Tags" ).
63- Order ("artifacts.created_at DESC" ).
86+ Order ("artifacts.created_at DESC" ). // Always pick the most recent
6487 First (
6588 & artifact ,
6689 & models.Artifact {ArtifactKey : in },
@@ -85,7 +108,27 @@ func (h *artifactRepo) Get(ctx context.Context, in models.ArtifactKey) (models.A
85108 return artifact , nil
86109}
87110
88- func (h * artifactRepo ) List (ctx context.Context , datasetKey models.DatasetKey , in models.ListModelsInput ) ([]models.Artifact , error ) {
111+ // Gets the artifact and locks the row
112+ func (h * artifactRepo ) getForUpdate (ctx context.Context , in models.ArtifactKey , db * gorm.DB ) (models.Artifact , error ) {
113+
114+ var artifact models.Artifact
115+ result := db .
116+ Clauses (clause.Locking {Strength : "UPDATE" }).
117+ Preload ("ArtifactData" ).
118+ Preload ("Partitions" , func (db * gorm.DB ) * gorm.DB {
119+ return db .WithContext (ctx ).Order ("partitions.created_at ASC" ) // preserve the order in which the partitions were created
120+ }).
121+ Preload ("Tags" ).
122+ Order ("artifacts.created_at DESC" ). // Always pick the most recent
123+ First (
124+ & artifact ,
125+ & models.Artifact {ArtifactKey : in },
126+ )
127+
128+ return artifact , result .Error
129+ }
130+
131+ func (h * artifactRepo ) ListAndFilterExpired (ctx context.Context , datasetKey models.DatasetKey , in models.ListModelsInput ) ([]models.Artifact , error ) {
89132 timer := h .repoMetrics .ListDuration .Start (ctx )
90133 defer timer .Stop ()
91134
0 commit comments