Skip to content

Commit

Permalink
Migrate to using new docPath for Firestore including projectID
Browse files Browse the repository at this point in the history 
fixes #213
  • Loading branch information
@nielm
nielm committed Feb 23, 2024
1 parent f6fcb49 commit 3801bf4
Show file tree
Hide file tree
Showing 2 changed files with 184 additions and 75 deletions.
37 changes: 32 additions & 5 deletions src/scaler/scaler-core/state.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -263,9 +263,8 @@ class StateFirestore extends State {
*/
get docRef() {
if (this._docRef == null) {
this.firestore = new firestore.Firestore({projectId: this.projectId});
this._docRef =
this.firestore.collection('spannerAutoscaler').doc(this.instanceId);
this._docRef = this.firestore
.doc(`spannerAutoscaler/state/${this.getSpannerId()}`);
}
return this._docRef;
}
Expand Down Expand Up @@ -299,9 +298,14 @@ class StateFirestore extends State {

/** @inheritdoc */
async get() {
const snapshot = await this.docRef.get(); // returns QueryDocumentSnapshot
let snapshot = await this.docRef.get(); // returns QueryDocumentSnapshot

if (!snapshot.exists) {
// It is possible that an old state doc exists in an old docref path...
snapshot = await this.checkAndReplaceOldDocRef();
}

if (!snapshot?.exists) {
this.data = await this.init();
} else {
this.data = snapshot.data();
Expand All @@ -310,7 +314,30 @@ class StateFirestore extends State {
return this.toMillis(this.data);
}

/** @inheritdoc */
/**
* Due to [issue 213](https://github.com/cloudspannerecosystem/autoscaler/issues/213)
* the docRef had to be changed, so check for an old doc at the old docref
* If it exists, copy it to the new docref, delete it and return it.
*/
async checkAndReplaceOldDocRef() {
try {
const oldDocRef =
this.firestore.doc(`spannerAutoscaler/${this.instanceId}`);
const snapshot = await oldDocRef.get();
if (snapshot.exists) {
await this.docRef.set(snapshot.data());
await oldDocRef.delete();
}
return snapshot;
} catch (e) {
// ignore
}
return null;
}

/**
* Update scaling timestamp in storage
*/
async set() {
await this.get(); // make sure doc exists

Expand Down
222 changes: 152 additions & 70 deletions src/scaler/scaler-core/test/state.test.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -72,10 +72,12 @@ afterEach(() => {
const DUMMY_TIMESTAMP = 1704110400000;

describe('stateFirestoreTests', () => {
let stubFirestoreInstance = sinon.createStubInstance(firestore.Firestore);
let collectionRef = sinon
.createStubInstance(firestore.CollectionReference);
let docRef = sinon.createStubInstance(firestore.DocumentReference);
/** @tyoe {sinon.SinonStubbedInstance<firestore.Firestore>} */
let stubFirestoreInstance;
/** @type {sinon.SinonStubbedInstance<firestore.DocumentReference<any>>} */
let newDocRef;
/** @type {sinon.SinonStubbedInstance<firestore.DocumentReference<any>>} */
let oldDocRef;

const autoscalerConfig = {
projectId: 'myProject',
Expand All @@ -86,14 +88,45 @@ describe('stateFirestoreTests', () => {
const DUMMY_FIRESTORE_TIMESTAMP = firestore.Timestamp
.fromMillis(DUMMY_TIMESTAMP);

const NEW_DOC_PATH =
'spannerAutoscaler/state/projects/myProject/instances/myInstance';
const OLD_DOC_PATH =
'spannerAutoscaler/myInstance';

/** @type {firestore.DocumentSnapshot<any>} */
// @ts-ignore
const EXISTING_DOC = {
exists: true,
data: () => {
return {
createdOn: DUMMY_FIRESTORE_TIMESTAMP,
updatedOn: DUMMY_FIRESTORE_TIMESTAMP,
lastScalingTimestamp:
DUMMY_FIRESTORE_TIMESTAMP,
};
},
};

/** @type {firestore.DocumentSnapshot<any>} */
// @ts-ignore
const NON_EXISTING_DOC = {
exists: false,
data: () => null,
};


beforeEach(() => {
// stub instances need to be recreated before each test.
stubFirestoreInstance = sinon.createStubInstance(firestore.Firestore);
collectionRef = sinon
.createStubInstance(firestore.CollectionReference);
docRef = sinon.createStubInstance(firestore.DocumentReference);
stubFirestoreConstructor.returns(stubFirestoreInstance);
stubFirestoreInstance.collection.returns(collectionRef);
collectionRef.doc.returns(docRef);
newDocRef = sinon.createStubInstance(firestore.DocumentReference);
oldDocRef = sinon.createStubInstance(firestore.DocumentReference);
stubFirestoreInstance.doc
.withArgs(NEW_DOC_PATH)
.returns(newDocRef);
stubFirestoreInstance.doc
.withArgs(OLD_DOC_PATH)
.returns(oldDocRef);
});

it('should create a StateFirestore object on spanner projectId',
Expand All @@ -119,35 +152,28 @@ describe('stateFirestoreTests', () => {
it('get() should read document from collection when exists',
async function() {
// @ts-ignore
docRef.get.returns(Promise.resolve({
exists: true,
data: () => {
return {
updatedOn: DUMMY_FIRESTORE_TIMESTAMP,
lastScalingTimestamp:
DUMMY_FIRESTORE_TIMESTAMP,
};
},
}));
newDocRef.get.returns(Promise.resolve(EXISTING_DOC));

const state = State.buildFor(autoscalerConfig);
const data = await state.get();

sinon.assert.calledWith(stubFirestoreInstance.collection,
'spannerAutoscaler');
sinon.assert.calledWith(collectionRef.doc, 'myInstance');
sinon.assert.calledOnce(newDocRef.get);
sinon.assert.calledWith(stubFirestoreInstance.doc, NEW_DOC_PATH);

// timestamp was converted...
assert.equals(data, {
updatedOn: DUMMY_TIMESTAMP,
lastScalingTimestamp: DUMMY_TIMESTAMP});
assert.equals(data,
{
createdOn: DUMMY_TIMESTAMP,
updatedOn: DUMMY_TIMESTAMP,
lastScalingTimestamp: DUMMY_TIMESTAMP,
});
});

it('get() should create a document when it does not exist',
async function() {
// @ts-ignore
docRef.get.returns(Promise.resolve({
exists: false,
}));
newDocRef.get.returns(Promise.resolve(NON_EXISTING_DOC));
oldDocRef.get.returns(Promise.resolve(NON_EXISTING_DOC));


const state = State.buildFor(autoscalerConfig);
const data = await state.get();
Expand All @@ -157,38 +183,94 @@ describe('stateFirestoreTests', () => {
createdOn: firestore.FieldValue.serverTimestamp(),
};

assert.equals(docRef.set.getCall(0).args[0], expected);
sinon.assert.calledTwice(stubFirestoreInstance.doc);
// first call to create docref is for the "new" path
assert.equals(stubFirestoreInstance.doc.getCall(0).args[0],
NEW_DOC_PATH);
// second call to create docref is for the "old" path
assert.equals(stubFirestoreInstance.doc.getCall(1).args[0],
OLD_DOC_PATH);

sinon.assert.calledOnce(newDocRef.get);
sinon.assert.calledOnce(oldDocRef.get);

sinon.assert.calledOnce(newDocRef.set);
assert.equals(newDocRef.set.getCall(0).args[0], expected);
assert.equals(data, expected);
});

it('get() should copy document from old location to new if missing in new',
async function() {
/**
* Due to [issue 213](https://github.com/cloudspannerecosystem/autoscaler/issues/213)
* the docRef had to be changed, so check for an old doc at the old
* docref, if it exists, copy it to the new docref, delete it and
* return it.
*/
newDocRef.get.returns(Promise.resolve(NON_EXISTING_DOC));
oldDocRef.get.returns(Promise.resolve(EXISTING_DOC));

const state = State.buildFor(autoscalerConfig);
const data = await state.get();

// Expected value set and returned is the old doc.
const expected = {
lastScalingTimestamp: DUMMY_TIMESTAMP,
createdOn: DUMMY_TIMESTAMP,
updatedOn: DUMMY_TIMESTAMP,
};

sinon.assert.calledTwice(stubFirestoreInstance.doc);
// first call to create docref is for the "new" path
assert.equals(stubFirestoreInstance.doc.getCall(0).args[0],
NEW_DOC_PATH);
// second call to create docref is for the "old" path
assert.equals(stubFirestoreInstance.doc.getCall(1).args[0],
OLD_DOC_PATH);

sinon.assert.calledOnce(newDocRef.get);
sinon.assert.calledOnce(oldDocRef.get);

// Copy data from existing doc in old location to new location.
sinon.assert.calledOnce(newDocRef.set);
assert.equals(newDocRef.set.getCall(0).args[0], EXISTING_DOC.data());
sinon.assert.calledOnce(oldDocRef.delete);

// return data from existing doc.
assert.equals(data, expected);
});


it('set() should write document to collection',
async function() {
// set calls get(), so give it a doc to return...
// @ts-ignore
docRef.get.returns(Promise.resolve({
exists: true,
data: () => {
return {
test: 'testdoc',
};
},
}));
newDocRef.get.returns(Promise.resolve(EXISTING_DOC));

const state = State.buildFor(autoscalerConfig);
await state.set();

assert.equals(docRef.update.getCall(0).args[0], {
sinon.assert.calledOnce(stubFirestoreInstance.doc);
assert.equals(stubFirestoreInstance.doc.getCall(0).args[0],
NEW_DOC_PATH);

sinon.assert.calledOnce(newDocRef.update);
assert.equals(newDocRef.update.getCall(0).args[0], {
updatedOn: firestore.FieldValue.serverTimestamp(),
lastScalingTimestamp: firestore.FieldValue.serverTimestamp(),
});
});
});


describe('stateSpannerTests', () => {
let stubSpannerClient = sinon.createStubInstance(spanner.Spanner);
let stubSpannerInstance = sinon.createStubInstance(spanner.Instance);
let stubSpannerDatabase = sinon.createStubInstance(spanner.Database);
let stubSpannerTable = sinon.createStubInstance(spanner.Table);
/** @type {sinon.SinonStubbedInstance<spanner.Spanner>} */
let stubSpannerClient;
/** @type {sinon.SinonStubbedInstance<spanner.Instance>} */
let stubSpannerInstance;
/** @type {sinon.SinonStubbedInstance<spanner.Database>} */
let stubSpannerDatabase;
/** @type {sinon.SinonStubbedInstance<spanner.Table>} */
let stubSpannerTable;

const autoscalerConfig = {
projectId: 'myProject',
Expand All @@ -213,6 +295,15 @@ describe('stateSpannerTests', () => {
},
};

const VALID_ROW = {
toJSON: () => {
return {
lastScalingTimestamp: new Date(DUMMY_TIMESTAMP),
createdOn: new Date(DUMMY_TIMESTAMP),
};
},
};

const DUMMY_SPANNER_ISO_TIME = spanner.Spanner.timestamp(DUMMY_TIMESTAMP)
.toISOString();

Expand All @@ -223,13 +314,19 @@ describe('stateSpannerTests', () => {
stubSpannerTable = sinon.createStubInstance(spanner.Table);

stubSpannerConstructor.returns(stubSpannerClient);
stubSpannerClient.instance.returns(stubSpannerInstance);
stubSpannerInstance.database.returns(stubSpannerDatabase);
stubSpannerDatabase.table.returns(stubSpannerTable);
stubSpannerClient.instance
.withArgs(autoscalerConfig.stateDatabase.instanceId)
.returns(stubSpannerInstance);
stubSpannerInstance.database
.withArgs(autoscalerConfig.stateDatabase.databaseId)
.returns(stubSpannerDatabase);
stubSpannerDatabase.table
.withArgs('spannerAutoscaler')
.returns(stubSpannerTable);
});


it('should create a StateSpanner object on spanner projectId',
it('should create a StateSpanner object connecting to spanner projectId',
function() {
const config = {
...autoscalerConfig,
Expand All @@ -238,7 +335,7 @@ describe('stateSpannerTests', () => {
const state = State.buildFor(config);
assert.equals(state.constructor.name, 'StateSpanner');
sinon.assert.calledWith(stubSpannerConstructor,
{projectId: 'myProject'});
{projectId: autoscalerConfig.projectId});
sinon.assert.calledWith(stubSpannerClient.instance,
autoscalerConfig.stateDatabase.instanceId);
sinon.assert.calledWith(stubSpannerInstance.database,
Expand All @@ -252,7 +349,7 @@ describe('stateSpannerTests', () => {
const state = State.buildFor(autoscalerConfig);
assert.equals(state.constructor.name, 'StateSpanner');
sinon.assert.calledWith(stubSpannerConstructor,
{projectId: 'stateProject'});
{projectId: autoscalerConfig.stateProjectId});
sinon.assert.calledWith(stubSpannerClient.instance,
autoscalerConfig.stateDatabase.instanceId);
sinon.assert.calledWith(stubSpannerInstance.database,
Expand All @@ -261,18 +358,10 @@ describe('stateSpannerTests', () => {
'spannerAutoscaler');
});

it('get() should read document from collection when exists',
it('get() should read document from table when exists',
async function() {
// @ts-ignore
stubSpannerTable.read.returns(Promise.resolve([[
{
toJSON: () => {
return {
lastScalingTimestamp: new Date(DUMMY_TIMESTAMP),
createdOn: new Date(DUMMY_TIMESTAMP),
};
},
}]]));
stubSpannerTable.read.returns(Promise.resolve([[VALID_ROW]]));

const state = State.buildFor(autoscalerConfig);
const data = await state.get();
Expand Down Expand Up @@ -309,21 +398,14 @@ describe('stateSpannerTests', () => {
});
});

it('set() should write document to collection',
it('set() should write document to table',
async function() {
// set calls get(), so give it a doc to return...
// @ts-ignore
stubSpannerTable.read.returns(Promise.resolve([[
{
toJSON: () => {
return {
lastScalingTimestamp: new Date(0),
createdOn: new Date(0),
};
},
}]]));
stubSpannerTable.read.returns(Promise.resolve([[VALID_ROW]]));

const state = State.buildFor(autoscalerConfig);

// make state.now return a fixed value
const nowfunc = sinon.stub();
sinon.replaceGetter(state, 'now', nowfunc);
Expand Down

0 comments on commit 3801bf4

Please sign in to comment.